test: use sqlite3 session in test_oauth_device_flow (#38782)

api/tests/unit_tests/services/test_oauth_device_flow.py

@@ -1,9 +1,14 @@
 from __future__ import annotations
 
 import uuid
+from datetime import UTC, datetime, timedelta
 from unittest.mock import MagicMock
 
+import pytest
+from sqlalchemy.orm import Session
+
 from libs.oauth_bearer import TOKEN_CACHE_KEY_FMT, AuthContext, SubjectType, TokenType
+from models.oauth import OAuthAccessToken
 from services.oauth_device_flow import (
     list_active_sessions,
     revoke_oauth_token,
@@ -11,13 +16,46 @@ from services.oauth_device_flow import (
     token_belongs_to_subject,
 )
 
+ACCOUNT_ID = uuid.UUID("11111111-1111-1111-1111-111111111111")
+OTHER_ACCOUNT_ID = uuid.UUID("22222222-2222-2222-2222-222222222222")
+TOKEN_ID = uuid.UUID("33333333-3333-3333-3333-333333333333")
+OTHER_TOKEN_ID = uuid.UUID("44444444-4444-4444-4444-444444444444")
 
-def _account_ctx() -> AuthContext:
+
+def _token(
+    *,
+    token_id: uuid.UUID = TOKEN_ID,
+    account_id: uuid.UUID | None = ACCOUNT_ID,
+    subject_email: str = "user@example.com",
+    subject_issuer: str = "dify:account",
+    token_hash: str | None = "live-hash",
+    expires_at: datetime | None = None,
+    revoked_at: datetime | None = None,
+    created_at: datetime | None = None,
+) -> OAuthAccessToken:
+    token = OAuthAccessToken(
+        subject_email=subject_email,
+        subject_issuer=subject_issuer,
+        account_id=str(account_id) if account_id is not None else None,
+        client_id="difyctl",
+        device_label="test-device",
+        prefix="dfoa_" if account_id is not None else "dfoe_",
+        token_hash=token_hash,
+        expires_at=expires_at or datetime.now(UTC) + timedelta(days=1),
+        revoked_at=revoked_at,
+    )
+    token.id = str(token_id)
+    if created_at is not None:
+        token.created_at = created_at
+    return token
+
+
+def _account_ctx(*, account_id: uuid.UUID = ACCOUNT_ID) -> AuthContext:
     return AuthContext(
         subject_type=SubjectType.ACCOUNT,
         subject_email="user@example.com",
         subject_issuer="dify:account",
-        account_id=uuid.uuid4(),
+        account_id=account_id,
         client_id="difyctl",
         scopes=frozenset({"full"}),
         token_id=uuid.uuid4(),
@@ -73,38 +111,47 @@ def test_subject_match_clauses_external_sso_requires_null_account_id():
 # ---------------------------------------------------------------------------
 
 
-def test_revoke_oauth_token_invalidates_redis_cache_when_live_hash_seen():
+@pytest.mark.parametrize("sqlite_session", [(OAuthAccessToken,)], indirect=True)
+def test_revoke_oauth_token_invalidates_redis_cache_when_live_hash_seen(sqlite_session: Session):
     """Happy path: snapshot finds a live ``token_hash`` → UPDATE runs +
     Redis cache entry is DEL'd so the next bearer probe re-reads the now
     revoked row from DB.
     """
-    session = MagicMock()
-    session.query.return_value.filter.return_value.one_or_none.return_value = ("live-hash",)
+    sqlite_session.add(_token())
+    sqlite_session.commit()
 
     redis = MagicMock()
 
-    revoke_oauth_token(redis, "token-id", session=session)
+    revoke_oauth_token(redis, str(TOKEN_ID), session=sqlite_session)
 
-    assert session.execute.called  # UPDATE ... WHERE revoked_at IS NULL
-    assert session.commit.called
+    assert not sqlite_session.in_transaction()
+    persisted = sqlite_session.get(OAuthAccessToken, str(TOKEN_ID))
+    assert persisted is not None
+    assert persisted.token_hash is None
+    assert persisted.revoked_at is not None
     redis.delete.assert_called_once_with(TOKEN_CACHE_KEY_FMT.format(hash="live-hash"))
 
 
-def test_revoke_oauth_token_is_idempotent_when_already_revoked():
+@pytest.mark.parametrize("sqlite_session", [(OAuthAccessToken,)], indirect=True)
+def test_revoke_oauth_token_is_idempotent_when_already_revoked(sqlite_session: Session):
     """Second call (or race-loser): no live hash → UPDATE still runs (it
     is itself idempotent thanks to ``WHERE revoked_at IS NULL``) but the
     Redis invalidation is skipped because there's no cache entry to
     drop.
     """
-    session = MagicMock()
-    session.query.return_value.filter.return_value.one_or_none.return_value = None
+    revoked_at = datetime.now(UTC) - timedelta(minutes=1)
+    sqlite_session.add(_token(token_hash=None, revoked_at=revoked_at))
+    sqlite_session.commit()
 
     redis = MagicMock()
 
-    revoke_oauth_token(redis, "token-id", session=session)
+    revoke_oauth_token(redis, str(TOKEN_ID), session=sqlite_session)
 
-    assert session.execute.called
-    assert session.commit.called
+    assert not sqlite_session.in_transaction()
+    persisted = sqlite_session.get(OAuthAccessToken, str(TOKEN_ID))
+    assert persisted is not None
+    assert persisted.token_hash is None
+    assert persisted.revoked_at is not None
     redis.delete.assert_not_called()
 
 
@@ -113,34 +160,60 @@ def test_revoke_oauth_token_is_idempotent_when_already_revoked():
 # ---------------------------------------------------------------------------
 
 
-def test_list_active_sessions_returns_session_execute_rows():
-    """Thin delegation: the helper materialises whatever
-    ``session.execute(...).scalars().all()`` returns into a list. The
-    ``.scalars()`` step unwraps each one-element ``Row`` so callers see
-    bare ``OAuthAccessToken`` entities (matches the declared return
-    type).
-    """
-    from datetime import UTC, datetime
+@pytest.mark.parametrize("sqlite_session", [(OAuthAccessToken,)], indirect=True)
+def test_list_active_sessions_returns_only_live_subject_tokens(sqlite_session: Session):
+    """Only live, hashed rows for the authenticated subject are returned newest-first."""
 
-    session = MagicMock()
-    fake_rows = [MagicMock(), MagicMock()]
-    session.execute.return_value.scalars.return_value.all.return_value = fake_rows
+    now = datetime.now(UTC)
+    active_new = _token(token_id=TOKEN_ID, created_at=now - timedelta(minutes=1))
+    active_old = _token(token_id=OTHER_TOKEN_ID, created_at=now - timedelta(minutes=2))
+    expired = _token(
+        token_id=uuid.UUID(int=5),
+        expires_at=now - timedelta(seconds=1),
+        created_at=now - timedelta(minutes=3),
+    )
+    revoked = _token(
+        token_id=uuid.UUID(int=6),
+        token_hash=None,
+        revoked_at=now - timedelta(seconds=1),
+        created_at=now - timedelta(minutes=4),
+    )
+    hashless = _token(
+        token_id=uuid.UUID(int=7),
+        token_hash=None,
+        created_at=now - timedelta(minutes=5),
+    )
+    other_account = _token(
+        token_id=uuid.UUID(int=8),
+        account_id=OTHER_ACCOUNT_ID,
+        created_at=now - timedelta(minutes=6),
+    )
+    external_sso = _token(
+        token_id=uuid.UUID(int=9),
+        account_id=None,
+        subject_email="user@example.com",
+        subject_issuer="https://idp.example.com",
+        created_at=now - timedelta(minutes=7),
+    )
+    sqlite_session.add_all([active_new, active_old, expired, revoked, hashless, other_account, external_sso])
+    sqlite_session.commit()
 
-    out = list_active_sessions(_account_ctx(), datetime.now(UTC), session=session)
+    out = list_active_sessions(_account_ctx(), now, session=sqlite_session)
 
-    assert out == fake_rows
-    assert session.execute.called
+    assert [token.id for token in out] == [str(TOKEN_ID), str(OTHER_TOKEN_ID)]
 
 
-def test_token_belongs_to_subject_true_when_row_present():
-    session = MagicMock()
-    session.execute.return_value.first.return_value = ("some-id",)
+@pytest.mark.parametrize("sqlite_session", [(OAuthAccessToken,)], indirect=True)
+def test_token_belongs_to_subject_true_when_row_present(sqlite_session: Session):
+    sqlite_session.add(_token())
+    sqlite_session.commit()
 
-    assert token_belongs_to_subject("token-id", _account_ctx(), session=session) is True
+    assert token_belongs_to_subject(str(TOKEN_ID), _account_ctx(), session=sqlite_session) is True
 
 
-def test_token_belongs_to_subject_false_when_no_row():
-    session = MagicMock()
-    session.execute.return_value.first.return_value = None
+@pytest.mark.parametrize("sqlite_session", [(OAuthAccessToken,)], indirect=True)
+def test_token_belongs_to_subject_false_for_other_account(sqlite_session: Session):
+    sqlite_session.add(_token(account_id=OTHER_ACCOUNT_ID))
+    sqlite_session.commit()
 
-    assert token_belongs_to_subject("token-id", _account_ctx(), session=session) is False
+    assert token_belongs_to_subject(str(TOKEN_ID), _account_ctx(), session=sqlite_session) is False