Jack1808/Claude_Code

1from unittest.mock import MagicMock, patch2 3from fastapi.testclient import TestClient4 5from api.app import app6from providers.nvidia_nim import NvidiaNimProvider7 8# Mock provider9mock_provider = MagicMock(spec=NvidiaNimProvider)10 11# Track stream_response calls for test_model_mapping12_stream_response_calls = []13 14 15async def _mock_stream_response(*args, **kwargs):16    """Minimal async generator for streaming tests."""17    _stream_response_calls.append((args, kwargs))18    yield "event: message_start\ndata: {}\n\n"19    yield "[DONE]\n\n"20 21 22mock_provider.stream_response = _mock_stream_response23 24# Patch get_provider_for_type to always return mock_provider25_patcher = patch("api.routes.get_provider_for_type", return_value=mock_provider)26_patcher.start()27 28client = TestClient(app)29 30 31def test_root():32    response = client.get("/")33    assert response.status_code == 20034    assert response.json()["status"] == "ok"35 36 37def test_health():38    response = client.get("/health")39    assert response.status_code == 20040    assert response.json()["status"] == "healthy"41 42 43def test_create_message_stream():44    """Create message returns streaming response."""45    payload = {46        "model": "claude-3-sonnet",47        "messages": [{"role": "user", "content": "Hi"}],48        "max_tokens": 100,49        "stream": True,50    }51    response = client.post("/v1/messages", json=payload)52    assert response.status_code == 20053    assert "text/event-stream" in response.headers.get("content-type", "")54    content = b"".join(response.iter_bytes())55    assert b"message_start" in content or b"event:" in content56 57 58def test_model_mapping():59    # Test Haiku mapping60    _stream_response_calls.clear()61    payload_haiku = {62        "model": "claude-3-haiku-20240307",63        "messages": [{"role": "user", "content": "Hi"}],64        "max_tokens": 100,65        "stream": True,66    }67    client.post("/v1/messages", json=payload_haiku)68    assert len(_stream_response_calls) == 169    args = _stream_response_calls[0][0]70    assert args[0].model != "claude-3-haiku-20240307"71    assert args[0].original_model == "claude-3-haiku-20240307"72 73 74def test_error_fallbacks():75    from providers.exceptions import (76        AuthenticationError,77        OverloadedError,78        RateLimitError,79    )80 81    base_payload = {82        "model": "test",83        "messages": [{"role": "user", "content": "Hi"}],84        "max_tokens": 10,85        "stream": True,86    }87 88    def _raise_auth(*args, **kwargs):89        raise AuthenticationError("Invalid Key")90 91    def _raise_rate_limit(*args, **kwargs):92        raise RateLimitError("Too Many Requests")93 94    def _raise_overloaded(*args, **kwargs):95        raise OverloadedError("Server Overloaded")96 97    # 1. Authentication Error (401)98    mock_provider.stream_response = _raise_auth99    response = client.post("/v1/messages", json=base_payload)100    assert response.status_code == 401101    assert response.json()["error"]["type"] == "authentication_error"102 103    # 2. Rate Limit (429)104    mock_provider.stream_response = _raise_rate_limit105    response = client.post("/v1/messages", json=base_payload)106    assert response.status_code == 429107    assert response.json()["error"]["type"] == "rate_limit_error"108 109    # 3. Overloaded (529)110    mock_provider.stream_response = _raise_overloaded111    response = client.post("/v1/messages", json=base_payload)112    assert response.status_code == 529113    assert response.json()["error"]["type"] == "overloaded_error"114 115    # Reset for subsequent tests116    mock_provider.stream_response = _mock_stream_response117 118 119def test_generic_exception_returns_500():120    """Non-ProviderError exceptions are caught and returned as HTTPException(500)."""121 122    def _raise_runtime(*args, **kwargs):123        raise RuntimeError("unexpected crash")124 125    mock_provider.stream_response = _raise_runtime126    response = client.post(127        "/v1/messages",128        json={129            "model": "test",130            "messages": [{"role": "user", "content": "Hi"}],131            "max_tokens": 10,132            "stream": True,133        },134    )135    assert response.status_code == 500136    mock_provider.stream_response = _mock_stream_response137 138 139def test_generic_exception_with_status_code():140    """Generic exception with status_code attribute uses that status (getattr fallback)."""141 142    class ExceptionWithStatus(RuntimeError):143        def __init__(self, msg: str, status_code: int = 500):144            super().__init__(msg)145            self.status_code = status_code146 147    def _raise_with_status(*args, **kwargs):148        raise ExceptionWithStatus("bad gateway", 502)149 150    mock_provider.stream_response = _raise_with_status151    response = client.post(152        "/v1/messages",153        json={154            "model": "test",155            "messages": [{"role": "user", "content": "Hi"}],156            "max_tokens": 10,157            "stream": True,158        },159    )160    assert response.status_code == 502161    mock_provider.stream_response = _mock_stream_response162 163 164def test_generic_exception_empty_message_returns_non_empty_detail():165    """Exceptions with empty __str__ still return a readable HTTP detail."""166 167    class SilentError(RuntimeError):168        def __str__(self):169            return ""170 171    def _raise_silent(*args, **kwargs):172        raise SilentError()173 174    mock_provider.stream_response = _raise_silent175    response = client.post(176        "/v1/messages",177        json={178            "model": "test",179            "messages": [{"role": "user", "content": "Hi"}],180            "max_tokens": 10,181            "stream": True,182        },183    )184    assert response.status_code == 500185    assert response.json()["detail"] != ""186    mock_provider.stream_response = _mock_stream_response187 188 189def test_count_tokens_endpoint():190    """count_tokens endpoint returns token count."""191    response = client.post(192        "/v1/messages/count_tokens",193        json={"model": "test", "messages": [{"role": "user", "content": "Hello"}]},194    )195    assert response.status_code == 200196    assert "input_tokens" in response.json()197 198 199def test_stop_endpoint_no_handler_no_cli_503():200    """POST /stop without handler or cli_manager returns 503."""201    # Ensure no handler or cli_manager on app state202    if hasattr(app.state, "message_handler"):203        delattr(app.state, "message_handler")204    if hasattr(app.state, "cli_manager"):205        delattr(app.state, "cli_manager")206    response = client.post("/stop")207    assert response.status_code == 503208