mirror of
https://18126008609:longquanjian123@gitee.com/feigong123/aurask.git
synced 2026-04-19 16:40:33 +00:00
78 lines
2.8 KiB
Python
78 lines
2.8 KiB
Python
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_LY_SSO_PASSWORD",
|
|
"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_ly_sso_login_creates_session_and_workspace(self) -> None:
|
|
os.environ["AURASK_LY_SSO_PASSWORD"] = "local-secret"
|
|
app = create_app(None, reset=True)
|
|
|
|
payload = app.login_with_ly_sso(username="ly-xujian1", password="local-secret")
|
|
context = app.auth.authenticate(f"Bearer {payload['token']}")
|
|
|
|
self.assertEqual(context["token_type"], "session")
|
|
self.assertEqual(payload["workspace"]["tenant_id"], payload["tenant"]["id"])
|
|
self.assertEqual(payload["quota"]["plan_code"], "free_trial")
|
|
self.assertEqual(len(app.list_workspaces(payload["tenant"]["id"])), 1)
|
|
|
|
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)
|
|
|
|
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()
|