aurask/tests/test_auth_sessions.py

from __future__ import annotations

import base64
import json
import os
import sys
import unittest
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "api"))

from aurask.app import create_app


class AuthSessionTests(unittest.TestCase):
    def setUp(self) -> None:
        self.env_keys = [
            "AURASK_GOOGLE_ENABLED",
            "AURASK_GOOGLE_ALLOW_UNVERIFIED_ID_TOKEN",
            "AURASK_GOOGLE_CLIENT_ID",
        ]
        self.old_env = {key: os.environ.get(key) for key in self.env_keys}

    def tearDown(self) -> None:
        for key, value in self.old_env.items():
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value

    def test_google_first_login_reuses_workspace_after_registration(self) -> None:
        os.environ["AURASK_GOOGLE_ENABLED"] = "true"
        os.environ["AURASK_GOOGLE_ALLOW_UNVERIFIED_ID_TOKEN"] = "true"
        os.environ["AURASK_GOOGLE_CLIENT_ID"] = "local-google-client"
        app = create_app(None, reset=True)
        token = self._fake_google_token(
            {
                "sub": "google-user-1",
                "email": "founder@example.com",
                "email_verified": True,
                "name": "Aurask Founder",
            }
        )

        first_login = app.login_with_google(id_token=token)
        second_login = app.login_with_google(id_token=token)

        self.assertTrue(first_login["is_new_user"])
        self.assertFalse(second_login["is_new_user"])
        self.assertEqual(first_login["workspace"]["id"], second_login["workspace"]["id"])
        self.assertEqual(len(app.list_workspaces(first_login["tenant"]["id"])), 1)
        self.assertEqual(app.auth.authenticate(f"Bearer {first_login['token']}")["token_type"], "session")

    def _fake_google_token(self, claims: dict) -> str:
        header = self._base64url({"alg": "none", "typ": "JWT"})
        payload = self._base64url(claims)
        return f"{header}.{payload}.signature"

    def _base64url(self, payload: dict) -> str:
        encoded = base64.urlsafe_b64encode(json.dumps(payload).encode("utf-8")).decode("utf-8")
        return encoded.rstrip("=")


if __name__ == "__main__":
    unittest.main()