test/test_openai_v1_models_pool.py

from __future__ import annotations

import unittest

from services.protocol import openai_v1_models


class _FakeAccountService:
    def __init__(self, token: str) -> None:
        self.token = token

    def peek_text_access_token(self) -> str:
        return self.token


class ModelListPoolTests(unittest.TestCase):
    def test_list_models_uses_pool_token_when_available(self) -> None:
        calls: list[str] = []

        class FakeBackend:
            def __init__(self, access_token: str = "") -> None:
                calls.append(access_token)

            def list_models(self) -> dict:
                return {"data": [{"id": "gpt-5", "object": "model"}]}

        original_backend = openai_v1_models.OpenAIBackendAPI
        original_account_service = openai_v1_models.account_service
        try:
            openai_v1_models.OpenAIBackendAPI = FakeBackend
            openai_v1_models.account_service = _FakeAccountService("pool-token")

            result = openai_v1_models.list_models()
        finally:
            openai_v1_models.OpenAIBackendAPI = original_backend
            openai_v1_models.account_service = original_account_service

        self.assertEqual(calls, ["pool-token"])
        model_ids = {item["id"] for item in result["data"]}
        self.assertIn("gpt-5", model_ids)

    def test_list_models_falls_back_to_anonymous_without_pool_token(self) -> None:
        calls: list[str] = []

        class FakeBackend:
            def __init__(self, access_token: str = "") -> None:
                calls.append(access_token)

            def list_models(self) -> dict:
                return {"data": [{"id": "auto", "object": "model"}]}

        original_backend = openai_v1_models.OpenAIBackendAPI
        original_account_service = openai_v1_models.account_service
        try:
            openai_v1_models.OpenAIBackendAPI = FakeBackend
            openai_v1_models.account_service = _FakeAccountService("")

            result = openai_v1_models.list_models()
        finally:
            openai_v1_models.OpenAIBackendAPI = original_backend
            openai_v1_models.account_service = original_account_service

        self.assertEqual(calls, [""])
        model_ids = {item["id"] for item in result["data"]}
        self.assertIn("auto", model_ids)

    def test_list_models_falls_back_to_anonymous_when_pool_probe_fails(self) -> None:
        calls: list[str] = []

        class FakeBackend:
            def __init__(self, access_token: str = "") -> None:
                self.access_token = access_token
                calls.append(access_token)

            def list_models(self) -> dict:
                if self.access_token:
                    raise RuntimeError("auth_models: HTTP 401")
                return {"data": [{"id": "auto", "object": "model"}]}

        original_backend = openai_v1_models.OpenAIBackendAPI
        original_account_service = openai_v1_models.account_service
        try:
            openai_v1_models.OpenAIBackendAPI = FakeBackend
            openai_v1_models.account_service = _FakeAccountService("bad-pool-token")

            result = openai_v1_models.list_models()
        finally:
            openai_v1_models.OpenAIBackendAPI = original_backend
            openai_v1_models.account_service = original_account_service

        self.assertEqual(calls, ["bad-pool-token", ""])
        model_ids = {item["id"] for item in result["data"]}
        self.assertIn("auto", model_ids)

    def test_list_models_uses_local_catalog_when_remote_model_discovery_fails(self) -> None:
        calls: list[str] = []

        class FakeBackend:
            def __init__(self, access_token: str = "") -> None:
                self.access_token = access_token
                calls.append(access_token)

            def list_models(self) -> dict:
                raise RuntimeError("bootstrap failed: status=403")

        original_backend = openai_v1_models.OpenAIBackendAPI
        original_account_service = openai_v1_models.account_service
        try:
            openai_v1_models.OpenAIBackendAPI = FakeBackend
            openai_v1_models.account_service = _FakeAccountService("blocked-pool-token")

            result = openai_v1_models.list_models()
        finally:
            openai_v1_models.OpenAIBackendAPI = original_backend
            openai_v1_models.account_service = original_account_service

        self.assertEqual(calls, ["blocked-pool-token", ""])
        model_ids = {item["id"] for item in result["data"]}
        self.assertIn("auto", model_ids)
        self.assertIn("gpt-5", model_ids)
        self.assertIn("gpt-image-2", model_ids)
        self.assertIn("codex-gpt-image-2", model_ids)

    def test_list_models_keeps_local_image_model_aliases(self) -> None:
        class FakeBackend:
            def __init__(self, access_token: str = "") -> None:
                pass

            def list_models(self) -> dict:
                return {"data": [{"id": "auto", "object": "model"}]}

        original_backend = openai_v1_models.OpenAIBackendAPI
        original_account_service = openai_v1_models.account_service
        try:
            openai_v1_models.OpenAIBackendAPI = FakeBackend
            openai_v1_models.account_service = _FakeAccountService("pool-token")

            result = openai_v1_models.list_models()
        finally:
            openai_v1_models.OpenAIBackendAPI = original_backend
            openai_v1_models.account_service = original_account_service

        model_ids = {item["id"] for item in result["data"]}
        self.assertIn("gpt-image-2", model_ids)
        self.assertIn("codex-gpt-image-2", model_ids)


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