diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 0b8e5878..e4c73991 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -23,7 +23,7 @@ jobs: python-version: '3.11' - name: Install dependencies - run: pip install httpx==0.28.1 pytest pytest-asyncio anyio pyyaml fastapi "pydantic>=2.0,<3.0" uvicorn python-multipart asyncpg langfuse redis + run: pip install httpx==0.28.1 pytest pytest-asyncio anyio pyyaml fastapi "pydantic>=2.0,<3.0" uvicorn python-multipart asyncpg langfuse redis aiofiles - name: Run Unit Tests run: CONFIG_PATH=router/config.yaml PYTHONPATH=.:router pytest --ignore=test_agy_behavior.py --ignore=test_agy_tiers.py diff --git a/router/Dockerfile b/router/Dockerfile index cd0f3653..49d9c6a6 100644 --- a/router/Dockerfile +++ b/router/Dockerfile @@ -3,7 +3,7 @@ FROM python:3.14-slim WORKDIR /app # Install deps in a single layer (no pip cache, no extra files) -RUN pip install --no-cache-dir fastapi "pydantic>=2.0,<3.0" uvicorn httpx pyyaml python-multipart asyncpg langfuse redis +RUN pip install --no-cache-dir fastapi "pydantic>=2.0,<3.0" uvicorn httpx pyyaml python-multipart asyncpg langfuse redis aiofiles # Copy all source in one layer — removes dead config COPY (volume-mounted at runtime) COPY main.py agy_proxy.py circuit_breaker.py aa_scores.json free_models_roster.json /app/ diff --git a/router/main.py b/router/main.py index c5ec94f9..f5bc30a9 100644 --- a/router/main.py +++ b/router/main.py @@ -1,4 +1,5 @@ import os +import aiofiles import re import sys import json @@ -3384,21 +3385,23 @@ class AnnotationItem(BaseModel): _annotations_cache = {} -def _read_annotations_sync(path) -> dict: +async def _read_annotations_async(path) -> dict: import copy # Do not swallow OSError if file doesn't exist to preserve original behavior. # The caller (save_annotations) handles the exception when reading existing annotations. - current_mtime = os.path.getmtime(path) + current_mtime = await asyncio.to_thread(os.path.getmtime, path) cache_entry = _annotations_cache.get(path) if cache_entry is None or current_mtime != cache_entry["mtime"]: - with open(path, "r", encoding="utf-8") as f: - data = json.load(f) + async with aiofiles.open(path, "r", encoding="utf-8") as f: + # Read asynchronously, but parse in a thread pool to avoid blocking event loop + content = await f.read() + data = await asyncio.to_thread(json.loads, content) _annotations_cache[path] = {"mtime": current_mtime, "data": data} - return copy.deepcopy(_annotations_cache[path]["data"]) + return await asyncio.to_thread(copy.deepcopy, _annotations_cache[path]["data"]) @app.post("/dashboard/save-annotations") async def save_annotations(payload: Dict[str, AnnotationItem]): @@ -3451,7 +3454,7 @@ async def save_annotations(payload: Dict[str, AnnotationItem]): async with annotations_lock: if ann_path.exists(): try: - existing = await asyncio.to_thread(_read_annotations_sync, str(ann_path)) + existing = await _read_annotations_async(str(ann_path)) except Exception as read_err: logger.warning(f"Could not read existing annotations: {read_err}. Overwriting.") diff --git a/scripts/README.md b/scripts/README.md index 0139e8d3..6d05eeac 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -79,6 +79,9 @@ The integration test suite is located in the root directory. Tests are categoriz ### Performance & Monitoring Tests - **`test_stream_latency.py`**: Measures Time-To-First-Token (TTFT) and token generation speed. +### Dashboard & Annotations Tests +- **`tests/test_read_annotations_async.py`**: Unit tests for the asynchronous dashboard annotation reading and caching mechanism. + ### Simulation Tests - **`test_agy_behavior.py`**: Asserts the behavior of the `agy` CLI client under quota limits. - **`test_quota_reset.sh`**: Simulates/triggers quota reset conditions. diff --git a/start-stack.sh b/start-stack.sh index 381b864c..03bb1220 100755 --- a/start-stack.sh +++ b/start-stack.sh @@ -141,6 +141,17 @@ if [ -z "$LANGFUSE_INIT_USER_PASSWORD" ]; then echo "✓ Generated new LANGFUSE_INIT_USER_PASSWORD and saved to $ENV_FILE" fi +if [ -z "$REDIS_AUTH" ]; then + REDIS_AUTH="$(openssl rand -hex 16)" + echo "REDIS_AUTH=\"$REDIS_AUTH\"" >> "$ENV_FILE" + echo "✓ Generated new REDIS_AUTH and saved to $ENV_FILE" +fi + +if [ -z "$CLICKHOUSE_PASSWORD" ]; then + CLICKHOUSE_PASSWORD="$(openssl rand -hex 16)" + echo "CLICKHOUSE_PASSWORD=\"$CLICKHOUSE_PASSWORD\"" >> "$ENV_FILE" + echo "✓ Generated new CLICKHOUSE_PASSWORD and saved to $ENV_FILE" +fi if [ -z "$ROUTER_API_KEY" ]; then ROUTER_API_KEY="$(openssl rand -hex 32)" diff --git a/tests/test_read_annotations_sync.py b/tests/test_read_annotations_async.py similarity index 56% rename from tests/test_read_annotations_sync.py rename to tests/test_read_annotations_async.py index 9ecf8320..b401bc8b 100644 --- a/tests/test_read_annotations_sync.py +++ b/tests/test_read_annotations_async.py @@ -1,35 +1,44 @@ import pytest -from unittest.mock import patch, mock_open +from unittest.mock import patch, AsyncMock, MagicMock import copy - -# Mock config so router.main can be imported import sys import os +import json +import asyncio # Ensure the root directory is in the path sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) -# Need to set up environment or mock configuration depending on how it's done in the rest of tests -from router.main import _read_annotations_sync import router.main +from router.main import _read_annotations_async @pytest.fixture(autouse=True) def clear_cache(): router.main._annotations_cache.clear() yield -def test_read_annotations_sync_initial_read(): +@pytest.mark.asyncio +async def test_read_annotations_async_initial_read(): fake_path = "/tmp/annotations.json" fake_data = {"annotation1": "data1"} + # Mock aiofiles.open + mock_file = AsyncMock() + mock_file.read.return_value = '{"annotation1": "data1"}' + + mock_context_manager = MagicMock() + mock_context_manager.__aenter__ = AsyncMock(return_value=mock_file) + mock_context_manager.__aexit__ = AsyncMock(return_value=False) + + mock_aiofiles_open = MagicMock(return_value=mock_context_manager) + with patch("os.path.getmtime", return_value=100.0) as mock_getmtime, \ - patch("builtins.open", mock_open(read_data='{"annotation1": "data1"}')) as mock_file, \ - patch("json.load", return_value=fake_data) as mock_json_load: + patch("aiofiles.open", mock_aiofiles_open) as mock_open: - result = _read_annotations_sync(fake_path) + result = await _read_annotations_async(fake_path) mock_getmtime.assert_called_once_with(fake_path) - mock_file.assert_called_once_with(fake_path, "r", encoding="utf-8") + mock_open.assert_called_once_with(fake_path, "r", encoding="utf-8") assert result == fake_data # Verify cache is populated @@ -37,23 +46,28 @@ def test_read_annotations_sync_initial_read(): assert router.main._annotations_cache[fake_path]["mtime"] == 100.0 assert router.main._annotations_cache[fake_path]["data"] == fake_data -def test_read_annotations_sync_cache_hit(): +@pytest.mark.asyncio +async def test_read_annotations_async_cache_hit(): fake_path = "/tmp/annotations.json" fake_data = {"annotation1": "data1"} # Pre-populate cache router.main._annotations_cache[fake_path] = {"mtime": 100.0, "data": fake_data} + # Mock aiofiles.open (should NOT be called) + mock_aiofiles_open = MagicMock() + with patch("os.path.getmtime", return_value=100.0) as mock_getmtime, \ - patch("builtins.open", mock_open()) as mock_file: + patch("aiofiles.open", mock_aiofiles_open) as mock_open: - result = _read_annotations_sync(fake_path) + result = await _read_annotations_async(fake_path) mock_getmtime.assert_called_once_with(fake_path) - mock_file.assert_not_called() + mock_open.assert_not_called() assert result == fake_data -def test_read_annotations_sync_cache_invalidation(): +@pytest.mark.asyncio +async def test_read_annotations_async_cache_invalidation(): fake_path = "/tmp/annotations.json" fake_data_old = {"annotation1": "data1"} fake_data_new = {"annotation2": "data2"} @@ -61,21 +75,30 @@ def test_read_annotations_sync_cache_invalidation(): # Pre-populate cache with old mtime router.main._annotations_cache[fake_path] = {"mtime": 100.0, "data": fake_data_old} + mock_file = AsyncMock() + mock_file.read.return_value = '{"annotation2": "data2"}' + + mock_context_manager = MagicMock() + mock_context_manager.__aenter__ = AsyncMock(return_value=mock_file) + mock_context_manager.__aexit__ = AsyncMock(return_value=False) + + mock_aiofiles_open = MagicMock(return_value=mock_context_manager) + with patch("os.path.getmtime", return_value=200.0) as mock_getmtime, \ - patch("builtins.open", mock_open(read_data='{"annotation2": "data2"}')) as mock_file, \ - patch("json.load", return_value=fake_data_new) as mock_json_load: + patch("aiofiles.open", mock_aiofiles_open) as mock_open: - result = _read_annotations_sync(fake_path) + result = await _read_annotations_async(fake_path) mock_getmtime.assert_called_once_with(fake_path) - mock_file.assert_called_once_with(fake_path, "r", encoding="utf-8") + mock_open.assert_called_once_with(fake_path, "r", encoding="utf-8") assert result == fake_data_new # Verify cache is updated assert router.main._annotations_cache[fake_path]["mtime"] == 200.0 assert router.main._annotations_cache[fake_path]["data"] == fake_data_new -def test_read_annotations_sync_deepcopy(): +@pytest.mark.asyncio +async def test_read_annotations_async_deepcopy(): fake_path = "/tmp/annotations.json" fake_data = {"annotation1": {"nested": "value"}} @@ -84,21 +107,22 @@ def test_read_annotations_sync_deepcopy(): with patch("os.path.getmtime", return_value=100.0): # First read - result1 = _read_annotations_sync(fake_path) + result1 = await _read_annotations_async(fake_path) # Mutate the result result1["annotation1"]["nested"] = "mutated" # Second read - result2 = _read_annotations_sync(fake_path) + result2 = await _read_annotations_async(fake_path) # Verify second read returns original data, not mutated assert result2["annotation1"]["nested"] == "value" assert router.main._annotations_cache[fake_path]["data"]["annotation1"]["nested"] == "value" -def test_read_annotations_sync_file_not_found(): +@pytest.mark.asyncio +async def test_read_annotations_async_file_not_found(): fake_path = "/tmp/annotations.json" with patch("os.path.getmtime", side_effect=FileNotFoundError): with pytest.raises(FileNotFoundError): - _read_annotations_sync(fake_path) + await _read_annotations_async(fake_path)