Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import os | |
| import sys | |
| import unittest | |
| from pathlib import Path | |
| from unittest.mock import AsyncMock, patch | |
| import httpx | |
| SERVICE_ROOT = Path(__file__).resolve().parents[1] | |
| BACKEND_ROOT = SERVICE_ROOT.parent | |
| if str(SERVICE_ROOT) not in sys.path: | |
| sys.path.insert(0, str(SERVICE_ROOT)) | |
| from inference.client import build_payload, post # noqa: E402 | |
| from inference.config import reset_inference_manifest_cache # noqa: E402 | |
| from inference.resolve_credentials import resolve_api_key # noqa: E402 | |
| CANONICAL_MANIFEST = BACKEND_ROOT / "config" / "inference.manifest.json" | |
| class InferenceClientTests(unittest.IsolatedAsyncioTestCase): | |
| def setUp(self) -> None: | |
| reset_inference_manifest_cache() | |
| self._old_manifest_path = os.environ.pop("INFERENCE_MANIFEST_PATH", None) | |
| self._old_api_base = os.environ.pop("OPENAI_BASE_URL", None) | |
| self._old_timeout = os.environ.pop("TABULAR_AGENT_TIMEOUT_SECONDS", None) | |
| self._old_api_key = os.environ.pop("OPENAI_API_KEY", None) | |
| os.environ["INFERENCE_MANIFEST_PATH"] = str(CANONICAL_MANIFEST) | |
| os.environ["OPENAI_API_KEY"] = "test-key" | |
| os.environ["OPENAI_BASE_URL"] = "https://example.test/v1" | |
| os.environ["TABULAR_AGENT_TIMEOUT_SECONDS"] = "99" | |
| def tearDown(self) -> None: | |
| reset_inference_manifest_cache() | |
| for key, value in ( | |
| ("INFERENCE_MANIFEST_PATH", self._old_manifest_path), | |
| ("OPENAI_BASE_URL", self._old_api_base), | |
| ("TABULAR_AGENT_TIMEOUT_SECONDS", self._old_timeout), | |
| ("OPENAI_API_KEY", self._old_api_key), | |
| ): | |
| if value is not None: | |
| os.environ[key] = value | |
| else: | |
| os.environ.pop(key, None) | |
| async def test_post_uses_api_base_and_timeout(self) -> None: | |
| payload = build_payload( | |
| system_prompt="SYSTEM", | |
| user_prompt="USER", | |
| max_tokens=64, | |
| temperature=0.0, | |
| ) | |
| mock_response = httpx.Response( | |
| 200, | |
| json={ | |
| "choices": [{"message": {"content": "{\"ok\": true}"}}], | |
| "usage": {"total_tokens": 12}, | |
| }, | |
| request=httpx.Request("POST", "https://example.test"), | |
| ) | |
| mock_client = AsyncMock(spec=httpx.AsyncClient) | |
| mock_client.post = AsyncMock(return_value=mock_response) | |
| with patch("inference.client.httpx.AsyncClient", return_value=mock_client): | |
| text, usage, data = await post( | |
| payload, | |
| model="gemini-3-flash", | |
| feature_id="python-review", | |
| client=mock_client, | |
| ) | |
| self.assertEqual(text, "{\"ok\": true}") | |
| self.assertEqual(usage, {"total_tokens": 12}) | |
| self.assertIn("choices", data) | |
| called_url = mock_client.post.await_args.args[0] | |
| self.assertTrue( | |
| called_url == "https://example.test/v1/chat/completions" | |
| ) | |
| request = mock_client.post.await_args | |
| self.assertEqual( | |
| request.kwargs["headers"]["Authorization"], | |
| "Bearer test-key", | |
| ) | |
| self.assertEqual( | |
| request.kwargs["json"]["model"], | |
| "antigravity/gemini-3-flash", | |
| ) | |
| self.assertEqual(resolve_api_key("proxy-openai"), "test-key") | |
| def test_missing_api_key_returns_none(self) -> None: | |
| os.environ.pop("OPENAI_API_KEY", None) | |
| os.environ.pop("TABULAR_AGENT_API_KEY", None) | |
| os.environ.pop("GOOGLE_API_KEY", None) | |
| os.environ.pop("WORD_REVIEW_API_KEY", None) | |
| reset_inference_manifest_cache() | |
| self.assertIsNone(resolve_api_key("proxy-openai")) | |
| if __name__ == "__main__": | |
| unittest.main() | |