grok-register/tests/test_grok2api_remote_pool.py

import sys
import types
import unittest
from unittest.mock import patch

# Keep this unit test independent from optional browser/network dependencies.
drission = types.ModuleType("DrissionPage")
drission.Chromium = type("Chromium", (), {})
drission.ChromiumOptions = type("ChromiumOptions", (), {})
drission_errors = types.ModuleType("DrissionPage.errors")
drission_errors.PageDisconnectedError = type("PageDisconnectedError", (Exception,), {})
curl_cffi = types.ModuleType("curl_cffi")
curl_cffi.requests = types.SimpleNamespace()
sys.modules.setdefault("DrissionPage", drission)
sys.modules.setdefault("DrissionPage.errors", drission_errors)
sys.modules.setdefault("curl_cffi", curl_cffi)

import grok_register_ttk as app


class DummyResponse:
    def __init__(self, payload=None, status_code=200, reason="", headers=None, text=""):
        self._payload = payload or {}
        self.status_code = status_code
        self.reason = reason
        self.headers = headers or {}
        self.text = text

    def raise_for_status(self):
        if self.status_code >= 400:
            raise RuntimeError(f"HTTP Error {self.status_code}: {self.reason}")

    def json(self):
        return self._payload


class Grok2ApiRemotePoolTests(unittest.TestCase):
    def setUp(self):
        self.original_config = app.config.copy()

    def tearDown(self):
        app.config = self.original_config

    def _configure(self, **overrides):
        app.config.update({
            "grok2api_remote_base": "https://grok.example.com",
            "grok2api_remote_app_key": "app-secret",
            "grok2api_pool_name": "ssoBasic",
            "grok2api_allow_legacy_full_save": False,
            **overrides,
        })

    def test_remote_pool_falls_back_to_admin_api_prefix_when_root_tokens_add_is_404(self):
        self._configure()
        calls = []

        def fake_post(url, **kwargs):
            calls.append((url, kwargs))
            if url == "https://grok.example.com/tokens/add":
                return DummyResponse(status_code=404)
            return DummyResponse({"status": "success", "count": 1})

        with patch.object(app, "http_post", side_effect=fake_post):
            ok = app.add_token_to_grok2api_remote_pool("sso=abc123", email="a@example.com")

        self.assertTrue(ok)
        self.assertEqual([url for url, _ in calls], [
            "https://grok.example.com/tokens/add",
            "https://grok.example.com/admin/api/tokens/add",
        ])
        self.assertEqual(calls[-1][1]["params"], {"app_key": "app-secret"})
        self.assertEqual(calls[-1][1]["json"], {
            "tokens": ["abc123"],
            "pool": "basic",
            "tags": ["auto-register"],
        })

    def test_remote_pool_does_not_duplicate_admin_api_prefix_when_base_already_points_to_admin_api(self):
        self._configure(
            grok2api_remote_base="https://grok.example.com/admin/api",
            grok2api_pool_name="ssoSuper",
        )
        calls = []

        def fake_post(url, **kwargs):
            calls.append((url, kwargs))
            return DummyResponse({"status": "success", "count": 1})

        with patch.object(app, "http_post", side_effect=fake_post):
            ok = app.add_token_to_grok2api_remote_pool("sso=super123", email="a@example.com")

        self.assertTrue(ok)
        self.assertEqual([url for url, _ in calls], [
            "https://grok.example.com/admin/api/tokens/add",
        ])
        self.assertEqual(calls[0][1]["json"]["pool"], "super")

    def test_remote_pool_full_save_fallback_requires_opt_in_and_uses_etag(self):
        self._configure(grok2api_allow_legacy_full_save=True)
        get_calls = []
        post_calls = []

        def fake_post(url, **kwargs):
            post_calls.append((url, kwargs))
            if url.endswith("/tokens/add"):
                return DummyResponse(status_code=404)
            if url == "https://grok.example.com/admin/api/tokens":
                return DummyResponse({"status": "success"})
            return DummyResponse(status_code=404)

        def fake_get(url, **kwargs):
            get_calls.append((url, kwargs))
            if url == "https://grok.example.com/admin/api/tokens":
                return DummyResponse(
                    {"tokens": {"ssoBasic": []}},
                    headers={"ETag": '"version-7"'},
                )
            return DummyResponse(status_code=404)

        with patch.object(app, "http_post", side_effect=fake_post), \
                patch.object(app, "http_get", side_effect=fake_get):
            ok = app.add_token_to_grok2api_remote_pool("sso=fallback123", email="a@example.com")

        self.assertTrue(ok)
        self.assertEqual([url for url, _ in get_calls], [
            "https://grok.example.com/tokens",
            "https://grok.example.com/admin/api/tokens",
        ])
        self.assertEqual(post_calls[-1][0], "https://grok.example.com/admin/api/tokens")
        self.assertEqual(post_calls[-1][1]["headers"]["If-Match"], '"version-7"')
        self.assertEqual(post_calls[-1][1]["json"], {
            "ssoBasic": [{"token": "fallback123", "tags": ["auto-register"], "note": "a@example.com"}],
        })

    def test_remote_pool_legacy_fallback_is_disabled_by_default(self):
        self._configure()
        with patch.object(app, "http_post", return_value=DummyResponse(status_code=404)), \
                patch.object(app, "http_get") as get_mock:
            with self.assertRaises(app.RemoteTokenCompatibilityError):
                app.add_token_to_grok2api_remote_pool("abc")
        get_mock.assert_not_called()

    def test_remote_pool_500_does_not_fallback(self):
        self._configure(grok2api_allow_legacy_full_save=True)
        with patch.object(app, "http_post", return_value=DummyResponse(status_code=500, text="boom")), \
                patch.object(app, "http_get") as get_mock:
            with self.assertRaises(app.RemoteTokenRequestError):
                app.add_token_to_grok2api_remote_pool("abc")
        get_mock.assert_not_called()

    def test_remote_pool_legacy_fallback_rejects_missing_etag(self):
        self._configure(grok2api_allow_legacy_full_save=True)

        def fake_post(url, **kwargs):
            if url.endswith("/tokens/add"):
                return DummyResponse(status_code=404)
            return DummyResponse({"status": "success"})

        def fake_get(url, **kwargs):
            if url.endswith("/tokens"):
                return DummyResponse({"tokens": {"ssoBasic": []}})
            return DummyResponse(status_code=404)

        with patch.object(app, "http_post", side_effect=fake_post), \
                patch.object(app, "http_get", side_effect=fake_get):
            with self.assertRaises(app.RemoteTokenCompatibilityError):
                app.add_token_to_grok2api_remote_pool("abc")


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