Spaces:
Running
Running
| import pytest | |
| from fastapi.testclient import TestClient | |
| def override_env(monkeypatch): | |
| from contextlib import contextmanager | |
| from fastapi.testclient import TestClient | |
| import gateway | |
| def _override(auth_mode, api_key="valid-key"): | |
| monkeypatch.setenv("NEEDLE_AUTH_MODE", auth_mode) | |
| monkeypatch.setenv("NEEDLE_API_KEY", api_key) | |
| monkeypatch.setattr(gateway, "AUTH_MODE", auth_mode) | |
| monkeypatch.setattr(gateway, "API_KEY", api_key) | |
| monkeypatch.setattr(gateway, "_check_model_ready_sync", lambda: True) | |
| class MockAgent: | |
| def complete(self, text, max_new_tokens=512): | |
| return {"success": True, "function_calls": [], "suppressed_calls": [], "confidence": 0.9} | |
| def reset(self): | |
| pass | |
| monkeypatch.setattr(gateway, "_build_agent", lambda t, s: MockAgent()) | |
| gateway.scheduler.reset_for_test() | |
| gateway._inference_lock_obj = None | |
| with TestClient(gateway.app) as client: | |
| yield client | |
| return _override | |
| def test_auth_optional_no_token_readonly(override_env): | |
| with override_env("optional") as client: | |
| assert client.get("/health").status_code == 200 | |
| def test_auth_optional_valid_token_readonly(override_env): | |
| with override_env("optional") as client: | |
| assert client.get("/health", headers={"Authorization": "Bearer valid-key"}).status_code == 200 | |
| def test_auth_optional_invalid_token_readonly(override_env): | |
| with override_env("optional") as client: | |
| assert client.get("/health", headers={"Authorization": "Bearer invalid"}).status_code == 200 | |
| def test_auth_optional_no_token_routing(override_env): | |
| with override_env("optional") as client: | |
| res = client.post("/v1/route", json={"input": "test", "tools": [{"name": "test", "description": "desc", "parameters": {"type": "object"}}]}) | |
| assert res.status_code == 200 | |
| def test_auth_optional_valid_token_routing(override_env): | |
| with override_env("optional") as client: | |
| res = client.post("/v1/route", json={"input": "test", "tools": [{"name": "test", "description": "desc", "parameters": {"type": "object"}}]}, headers={"Authorization": "Bearer valid-key"}) | |
| assert res.status_code == 200 | |
| def test_auth_optional_no_token_mutation(override_env): | |
| with override_env("optional") as client: | |
| res = client.post("/v1/tools/register", json={"namespace": "ns", "tools": []}) | |
| assert res.status_code == 401 | |
| def test_auth_optional_valid_admin_token_mutation(override_env): | |
| with override_env("optional") as client: | |
| res = client.post("/v1/tools/register", json={"namespace": "ns", "tools": []}, headers={"X-API-Key": "valid-key"}) | |
| assert res.status_code == 200 | |
| def test_auth_required_no_token_readonly(override_env): | |
| with override_env("required") as client: | |
| assert client.get("/health").status_code == 401 | |
| def test_auth_required_valid_token_readonly(override_env): | |
| with override_env("required") as client: | |
| assert client.get("/health", headers={"Authorization": "Bearer valid-key"}).status_code == 200 | |
| def test_auth_required_invalid_token_readonly(override_env): | |
| with override_env("required") as client: | |
| assert client.get("/health", headers={"Authorization": "Bearer invalid"}).status_code == 401 | |
| def test_auth_required_no_token_routing(override_env): | |
| with override_env("required") as client: | |
| res = client.post("/v1/route", json={"input": "test", "tools": [{"name": "test", "description": "desc", "parameters": {"type": "object"}}]}) | |
| assert res.status_code == 401 | |
| def test_auth_required_valid_token_routing(override_env): | |
| with override_env("required") as client: | |
| res = client.post("/v1/route", json={"input": "test", "tools": [{"name": "test", "description": "desc", "parameters": {"type": "object"}}]}, headers={"Authorization": "Bearer valid-key"}) | |
| assert res.status_code == 200 | |
| def test_auth_required_no_token_mutation(override_env): | |
| with override_env("required") as client: | |
| res = client.post("/v1/tools/register", json={"namespace": "ns", "tools": []}) | |
| assert res.status_code == 401 | |
| def test_auth_required_valid_admin_token_mutation(override_env): | |
| with override_env("required") as client: | |
| res = client.post("/v1/tools/register", json={"namespace": "ns", "tools": []}, headers={"Authorization": "Bearer valid-key"}) | |
| assert res.status_code == 200 | |