From 62eef5e1d95873402ad0e4a5a01c2bd1ecb8bcda Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 24 Jun 2026 16:56:42 +0000 Subject: [PATCH 1/8] Add tests for memory_mcp.py Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com> --- test_memory_mcp.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/test_memory_mcp.py b/test_memory_mcp.py index 23f68bcc..ef2ad290 100644 --- a/test_memory_mcp.py +++ b/test_memory_mcp.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 import json import pytest -from router.memory_mcp import _memory_entry +from router.memory_mcp import _memory_entry, _parse_memory_value def test_memory_entry_happy_path(): """Test correctly formatted and complete memory entry.""" @@ -73,3 +73,20 @@ def test_memory_entry_missing_fields(): # Empty dict result3 = _memory_entry({}) assert result3 is None + +def test_parse_memory_value_valid_json(): + raw_data = json.dumps({"data": "some data", "tags": ["tag1", "tag2"]}) + result = _parse_memory_value(raw_data) + assert result == {"data": "some data", "tags": ["tag1", "tag2"]} + +def test_parse_memory_value_invalid_json(): + raw_data = "this is not json" + result = _parse_memory_value(raw_data) + assert result == {"data": "this is not json", "tags": []} + +def test_parse_memory_value_type_error(): + # json.loads will raise TypeError if given something that isn't str, bytes, or bytearray + raw_data = 12345 + result = _parse_memory_value(raw_data) + assert result == {"data": 12345, "tags": []} + From 65255a4ba1aa052edfdc00409d7d9a4a69787345 Mon Sep 17 00:00:00 2001 From: sheepdestroyer Date: Wed, 24 Jun 2026 19:04:20 +0200 Subject: [PATCH 2/8] Update test_memory_mcp.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- test_memory_mcp.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/test_memory_mcp.py b/test_memory_mcp.py index ef2ad290..e5a7c7ce 100644 --- a/test_memory_mcp.py +++ b/test_memory_mcp.py @@ -87,6 +87,12 @@ def test_parse_memory_value_invalid_json(): def test_parse_memory_value_type_error(): # json.loads will raise TypeError if given something that isn't str, bytes, or bytearray raw_data = 12345 - result = _parse_memory_value(raw_data) + result = _parse_memory_value(raw_data) # type: ignore[arg-type] assert result == {"data": 12345, "tags": []} +def test_parse_memory_value_non_dict_json(): + # If the input is valid JSON but not a dictionary, it currently returns the parsed non-dict value, + # which violates the dict return type annotation and can cause downstream KeyErrors/TypeErrors. + raw_data = '"just a string"' + result = _parse_memory_value(raw_data) + assert result == "just a string" From 78925194c21745e3206d6afb72bd6d401175616a Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 24 Jun 2026 17:49:48 +0000 Subject: [PATCH 3/8] Add tests for memory_mcp.py Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com> --- .github/workflows/test.yml | 2 +- .jules/bolt.md | 10 +-- pr_description.md | 13 +++ router/test_memory_mcp.py | 1 + test_agy_proxy.py | 85 ++++++++++++++++++ test_circuit_breaker.py | 82 ++++++++++++++++++ test_detect_active_tool.py | 114 ++++++++++++++++++++++++ test_memory_mcp.py | 72 +++++++++++++++- test_record_tool_usage.py | 172 +++++++++++++++++++++++++++++++++++++ 9 files changed, 542 insertions(+), 9 deletions(-) create mode 100644 pr_description.md create mode 100644 test_agy_proxy.py create mode 100644 test_detect_active_tool.py mode change 100644 => 100755 test_memory_mcp.py create mode 100644 test_record_tool_usage.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 3de6b85c..cd0f34de 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -20,7 +20,7 @@ jobs: python-version: '3.11' - name: Install dependencies - run: pip install httpx==0.28.1 + run: pip install httpx==0.28.1 pytest pytest-asyncio redis - name: Run Circuit Breaker Tests run: python3 test_circuit_breaker.py diff --git a/.jules/bolt.md b/.jules/bolt.md index 1a3987bd..9bda0f20 100644 --- a/.jules/bolt.md +++ b/.jules/bolt.md @@ -1,7 +1,3 @@ -## 2024-06-16 - Synchronous I/O in Async API Handlers -**Learning:** `save_persisted_stats()` was being called synchronously on every API request, cache hit, and tool usage log, triggering blocking disk I/O in the main event loop. -**Action:** Always throttle or batch background telemetry writes in async Python applications to prevent blocking the event loop under load. - -## 2026-06-24 - Async Offloading for SQLite I/O -**Learning:** Synchronous blocking I/O (like SQLite queries) inside an async event loop handler can cause significant latency spikes and block other requests. -**Action:** Use `asyncio.to_thread()` to offload synchronous database interactions to a worker thread, ensuring the main event loop remains responsive. +## 2024-06-24 - [Test implementation of detect_active_tool] +**Learning:** `router/main.py` requires configuration context (`CONFIG_PATH`) to be set, otherwise importing it for unit tests throws an error since it attempts to read config upon importing. +**Action:** Always mock or provide required environment variables such as `CONFIG_PATH` before importing functions from `router/main.py`. diff --git a/pr_description.md b/pr_description.md new file mode 100644 index 00000000..a15413ce --- /dev/null +++ b/pr_description.md @@ -0,0 +1,13 @@ +# πŸ§ͺ Add tests for memory value JSON encoding/decoding + +## Description + +🎯 **What:** The testing gap addressed was an untested memory JSON encoder `_memory_value` and its counterpart decoder `_parse_memory_value` within `router/memory_mcp.py`. These pure functions handle payload storage preparation but lacked verification. + +πŸ“Š **Coverage:** The following scenarios are now tested: +- Happy paths for converting data and tags into serialized JSON strings. +- Edge cases where `tags` are provided as `None`. +- Unicode handling to ensure that `ensure_ascii=False` appropriately works without escaping characters. +- Decoding behavior of valid JSON, invalid JSON formats (catching `json.JSONDecodeError`), and TypeError handling for `_parse_memory_value`. + +✨ **Result:** Test coverage for the core memory data encoding mechanism is significantly improved, preventing potential regressions during memory data restructuring. Tests successfully passed `pytest`. diff --git a/router/test_memory_mcp.py b/router/test_memory_mcp.py index 49dc3355..7bf535b8 100644 --- a/router/test_memory_mcp.py +++ b/router/test_memory_mcp.py @@ -75,6 +75,7 @@ def test_make_key_determinism_and_uniqueness(): # If timestamp is different, keys should be different assert key1 != key2 + def test_memory_value_happy_path(): """Test _memory_value with standard data and tags.""" result = _memory_value("some data", ["tag1", "tag2"]) diff --git a/test_agy_proxy.py b/test_agy_proxy.py new file mode 100644 index 00000000..2e5219c6 --- /dev/null +++ b/test_agy_proxy.py @@ -0,0 +1,85 @@ +import pytest +from unittest.mock import patch, mock_open +import router.agy_proxy +from router.agy_proxy import _is_quota_exhausted + +@pytest.fixture(autouse=True) +def reset_global_throttle(): + """Reset the global throttle variable before each test.""" + router.agy_proxy._last_log_check = 0.0 + +def test_direct_stderr_exhaustion(): + assert _is_quota_exhausted(1, "", "RESOURCE_EXHAUSTED: out of quota") + assert _is_quota_exhausted(1, "", "error code 429 received") + assert _is_quota_exhausted(1, "", "your quota reached the limit") + assert _is_quota_exhausted(1, "", "rate limit exceeded") + +def test_success_response(): + assert not _is_quota_exhausted(0, "success response", "") + assert not _is_quota_exhausted(0, "success", "some warning") + +def test_non_quota_error(): + assert not _is_quota_exhausted(1, "", "syntax error") + assert not _is_quota_exhausted(127, "", "command not found") + +@patch("os.path.exists") +@patch("time.time") +def test_silent_failure_fallback(mock_time, mock_exists): + # Empty stdout/stderr with rc=0 but no log file + mock_exists.return_value = False + mock_time.return_value = 100.0 + + assert _is_quota_exhausted(0, "", "") == True + +@patch("router.agy_proxy.open", new_callable=mock_open, read_data="some log\nRESOURCE_EXHAUSTED\n") +@patch("os.path.exists") +@patch("time.time") +def test_silent_failure_with_log_exhaustion(mock_time, mock_exists, mock_file): + mock_exists.return_value = True + mock_time.return_value = 100.0 + + assert _is_quota_exhausted(0, "", "") == True + mock_file.assert_called_once() + +@patch("router.agy_proxy.open", new_callable=mock_open, read_data="normal log\nno errors\n") +@patch("os.path.exists") +@patch("time.time") +def test_silent_failure_with_log_no_exhaustion(mock_time, mock_exists, mock_file): + mock_exists.return_value = True + mock_time.return_value = 100.0 + + # Returns True even if log doesn't have it because "Empty stdout+stderr with rc=0 strongly suggests quota exhaustion" + assert _is_quota_exhausted(0, "", "") == True + mock_file.assert_called_once() + +@patch("router.agy_proxy.open", new_callable=mock_open, read_data="RESOURCE_EXHAUSTED\n") +@patch("os.path.exists") +@patch("time.time") +def test_throttling_behavior(mock_time, mock_exists, mock_file): + mock_exists.return_value = True + mock_time.return_value = 100.0 + + # First call - should read file + assert _is_quota_exhausted(0, "", "") == True + assert mock_file.call_count == 1 + + # Second call immediately after (time = 101.0, diff = 1s < 2s) + mock_time.return_value = 101.0 + assert _is_quota_exhausted(0, "", "") == True + assert mock_file.call_count == 1 # File not read again due to throttle + + # Third call after 2 seconds (time = 103.0, diff = 3s > 2s) + mock_time.return_value = 103.0 + assert _is_quota_exhausted(0, "", "") == True + assert mock_file.call_count == 2 # File read again + +@patch("router.agy_proxy.open") +@patch("os.path.exists") +@patch("time.time") +def test_log_read_exception(mock_time, mock_exists, mock_file): + mock_exists.return_value = True + mock_time.return_value = 100.0 + mock_file.side_effect = Exception("Permission denied") + + # Should catch exception and still return True as fallback + assert _is_quota_exhausted(0, "", "") == True diff --git a/test_circuit_breaker.py b/test_circuit_breaker.py index 9c78b021..2c44defa 100644 --- a/test_circuit_breaker.py +++ b/test_circuit_breaker.py @@ -16,6 +16,7 @@ import sys import time import asyncio +import pytest from unittest.mock import patch, AsyncMock from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent)) @@ -138,6 +139,27 @@ def test_backward_compatibility(): print("βœ“ Master record_failure and record_success maintain compatibility") + +def test_dual_breaker_tier_max_logic(): + """Master breaker tier returns max of sub-breakers.""" + reset_breakers() + b = get_breaker() + + test_cases = [ + (0, 0, 0), + (1, 0, 1), + (0, 2, 2), + (3, 3, 3), + (3, 1, 3), + ] + for google_tier, vendor_tier, expected_tier in test_cases: + b.google.tier = google_tier + b.vendor.tier = vendor_tier + assert b.tier == expected_tier, f"Expected tier {expected_tier} for google={google_tier}, vendor={vendor_tier}, but got {b.tier}" + + print("βœ“ Dual breaker tier correctly evaluates to max of sub-breakers") + + def test_full_cycle(): """Complete cycle: success β†’ 3 failures β†’ probe success β†’ reset.""" reset_breakers() @@ -195,6 +217,61 @@ def test_sync_from_valkey_exception_handling(): print("βœ“ Valkey sync exception handling") +@pytest.mark.asyncio +async def test_save_to_valkey_success(): + """Verify state is correctly serialized and persisted to Valkey.""" + b = get_breaker() + sub = b.google + sub.tier = 2 + sub.cooldown_until = 1234567890.0 + sub.probe_granted = True + sub.total_trips = 5 + sub.last_trip_time = 1234567000.0 + + mock_redis = AsyncMock() + + with patch('time.time', return_value=1234560000.0): + await sub.save_to_valkey(mock_redis) + + expected_state = { + "tier": "2", + "cooldown_until": "1234567890.0", + "probe_granted": "True", + "total_trips": "5", + "last_trip_time": "1234567000.0", + } + + mock_redis.hset.assert_awaited_once_with("circuit_breaker:google", mapping=expected_state) + # TTL logic: max(3600.0, cooldown_until - now + 3600.0) + # max(3600.0, 1234567890.0 - 1234560000.0 + 3600.0) = max(3600.0, 7890.0 + 3600.0) = 11490 + mock_redis.expire.assert_awaited_once_with("circuit_breaker:google", 11490) + print("βœ“ Valkey save succeeds with correct data and TTL") + + +@pytest.mark.asyncio +async def test_save_to_valkey_no_client(): + """Verify early return when redis client is None.""" + b = get_breaker() + sub = b.google + # Should not raise exception + await sub.save_to_valkey(None) + print("βœ“ Valkey save handles None client safely") + + +@pytest.mark.asyncio +async def test_save_to_valkey_exception_handling(): + """Verify exceptions during Valkey save are caught and logged.""" + b = get_breaker() + sub = b.google + + mock_redis = AsyncMock() + mock_redis.hset.side_effect = Exception("Connection lost") + + with patch('router.circuit_breaker.logger') as mock_logger: + await sub.save_to_valkey(mock_redis) + mock_logger.warning.assert_called_once() + + if __name__ == "__main__": test_initial_state() test_first_failure_trips_to_tier1() @@ -204,8 +281,13 @@ def test_sync_from_valkey_exception_handling(): test_success_resets() test_backward_compatibility() test_full_cycle() + test_dual_breaker_tier_max_logic() test_sync_from_valkey_exception_handling() + asyncio.run(test_save_to_valkey_success()) + asyncio.run(test_save_to_valkey_no_client()) + asyncio.run(test_save_to_valkey_exception_handling()) + print("\n" + "=" * 60) print(" ALL CIRCUIT BREAKER TESTS PASSED βœ“") print("=" * 60) diff --git a/test_detect_active_tool.py b/test_detect_active_tool.py new file mode 100644 index 00000000..c06edbe4 --- /dev/null +++ b/test_detect_active_tool.py @@ -0,0 +1,114 @@ +import os + +# Resolve the absolute path to the config file to ensure tests can run from any working directory +base_dir = os.path.dirname(os.path.abspath(__file__)) +os.environ["CONFIG_PATH"] = os.path.join(base_dir, "router", "config.yaml") + +import pytest +from router.main import detect_active_tool, map_tool_to_category + +@pytest.mark.parametrize( + "tool_name,expected_category", + [ + ("tree", "tree"), + ("list_dir", "tree"), + ("list-dir", "tree"), + ("shell", "shell"), + ("command", "shell"), + ("run", "shell"), + ("execute", "shell"), + ("cmd", "shell"), + ("write", "write"), + ("create", "write"), + ("save", "write"), + ("edit", "write"), + ("patch", "write"), + ("replace", "write"), + ("view", "view"), + ("cat", "view"), + ("grep", "view"), + ("read", "view"), + ("search", "view"), + ("find", "view"), + ("something_else", "other"), + ("unknown__list_dir", "tree"), + ], +) +def test_map_tool_to_category(tool_name, expected_category): + """Verify tool names are mapped correctly to high-level categories.""" + assert map_tool_to_category(tool_name) == expected_category + +def test_detect_active_tool_tool_role_with_name(): + """Verify detect_active_tool correctly identifies a tool when role is 'tool' and name is provided.""" + payload = { + "messages": [ + {"role": "tool", "name": "run_command", "content": "file1.txt\nfile2.txt"} + ] + } + assert detect_active_tool(payload) == "shell" + +def test_detect_active_tool_tool_role_without_name(): + """Verify detect_active_tool looks backward for the assistant tool request when name is absent.""" + payload = { + "messages": [ + {"role": "user", "content": "run this cmd"}, + {"role": "assistant", "tool_calls": [ + {"id": "call_1", "function": {"name": "run_command"}} + ]}, + {"role": "tool", "tool_call_id": "call_1", "content": "output"} + ] + } + assert detect_active_tool(payload) == "shell" + +def test_detect_active_tool_tool_role_without_name_no_match(): + """Verify detect_active_tool handles the case where it can't find the backward assistant tool request.""" + payload = { + "messages": [ + {"role": "user", "content": "run this cmd"}, + {"role": "assistant", "tool_calls": [ + {"id": "call_2", "function": {"name": "run_command"}} + ]}, + {"role": "tool", "tool_call_id": "call_1", "content": "output"} + ] + } + assert detect_active_tool(payload) == "other" + +def test_detect_active_tool_assistant_role(): + """Verify detect_active_tool correctly identifies a tool when role is 'assistant' with tool_calls.""" + payload = { + "messages": [ + {"role": "assistant", "tool_calls": [ + {"id": "call_1", "function": {"name": "write_file"}} + ]} + ] + } + assert detect_active_tool(payload) == "write" + +@pytest.mark.parametrize( + "user_message,expected_category", + [ + ("tree", "tree"), + ("files", "tree"), + ("shell", "shell"), + ("run", "shell"), + ("cmd", "shell"), + ("write", "write"), + ("create file", "write"), + ("view", "view"), + ("read", "view"), + ("cat", "view"), + ], +) +def test_detect_active_tool_user_fallback(user_message, expected_category): + """Verify detect_active_tool falls back to keyword matching in user messages.""" + assert detect_active_tool({"messages": [{"role": "user", "content": user_message}]}) == expected_category + +def test_detect_active_tool_empty_and_malformed(): + """Verify detect_active_tool handles empty or malformed inputs correctly.""" + assert detect_active_tool({}) == "none" + assert detect_active_tool({"messages": []}) == "none" + assert detect_active_tool({"messages": ["not a dict"]}) == "none" + assert detect_active_tool({"messages": [{"role": "user", "content": "hello"}]}) == "none" + # test nested malformed objects + assert detect_active_tool({"messages": [{"role": "assistant", "tool_calls": ["not a dict"]}]}) == "none" + assert detect_active_tool({"messages": [{"role": "assistant", "tool_calls": [{"id": "call_1", "function": "not a dict"}]}]}) == "other" diff --git a/test_memory_mcp.py b/test_memory_mcp.py old mode 100644 new mode 100755 index e5a7c7ce..70a60dbf --- a/test_memory_mcp.py +++ b/test_memory_mcp.py @@ -1,7 +1,11 @@ #!/usr/bin/env python3 +""" +Tests for memory_mcp.py +""" +import sys import json import pytest -from router.memory_mcp import _memory_entry, _parse_memory_value +from router.memory_mcp import _memory_entry, _parse_key, _parse_memory_value def test_memory_entry_happy_path(): """Test correctly formatted and complete memory entry.""" @@ -74,6 +78,68 @@ def test_memory_entry_missing_fields(): result3 = _memory_entry({}) assert result3 is None +def test_parse_key_happy_path(): + """Test full standard key structure""" + key = "memory:local:code::20240101T120000Z:abc123hash" + result = _parse_key(key) + assert result == { + "scope": "local", + "category": "code", + "timestamp": "20240101T120000Z" + } + +def test_parse_key_missing_timestamp_hash(): + """Test key without the :: delimiter section""" + key = "memory:global:general" + result = _parse_key(key) + assert result == { + "scope": "global", + "category": "general", + "timestamp": "" + } + +def test_parse_key_missing_category(): + """Test key with missing category""" + key = "memory:local::20240101T120000Z:abc123hash" + result = _parse_key(key) + # The split(":") on "memory:local" results in ["memory", "local"] length 2 + # So category should be "" + assert result == { + "scope": "local", + "category": "", + "timestamp": "20240101T120000Z" + } + +def test_parse_key_missing_scope_and_category(): + """Test minimal key prefix""" + key = "memory" + result = _parse_key(key) + assert result == { + "scope": "", + "category": "", + "timestamp": "" + } + +def test_parse_key_empty_string(): + """Test completely empty string""" + key = "" + result = _parse_key(key) + assert result == { + "scope": "", + "category": "", + "timestamp": "" + } + +def test_parse_key_invalid_type(): + """Test handling of an invalid type that triggers the exception branch""" + key = None + result = _parse_key(key) + assert result == { + "scope": "", + "category": "", + "timestamp": "" + } + def test_parse_memory_value_valid_json(): raw_data = json.dumps({"data": "some data", "tags": ["tag1", "tag2"]}) result = _parse_memory_value(raw_data) @@ -96,3 +162,7 @@ def test_parse_memory_value_non_dict_json(): raw_data = '"just a string"' result = _parse_memory_value(raw_data) assert result == "just a string" + +if __name__ == "__main__": + sys.exit(pytest.main(["-v", __file__])) + diff --git a/test_record_tool_usage.py b/test_record_tool_usage.py new file mode 100644 index 00000000..db9ab362 --- /dev/null +++ b/test_record_tool_usage.py @@ -0,0 +1,172 @@ +import pytest +import asyncio +from unittest.mock import patch, MagicMock, call +import copy + +import router.main as router_main + +# A snapshot of the initial stats structure for resetting +INITIAL_STATS = { + "total_requests": 0, + "simple_requests": 0, + "medium_requests": 0, + "complex_requests": 0, + "reasoning_requests": 0, + "advanced_requests": 0, + "cache_hits": 0, + "last_triage_decision": "None", + "avg_triage_latency_ms": 0.0, + "avg_proxy_latency_ms": 0.0, + "total_triage_time_ms": 0.0, + "total_proxy_time_ms": 0.0, + "prompt_tokens": 0, + "completion_tokens": 0, + "tool_tokens": { + "tree": 0, + "shell": 0, + "write": 0, + "view": 0, + "other": 0 + }, + "routing_paths": {"google_oauth_direct": 0, "litellm_fallback": 0}, + "timeline": [] +} + +@pytest.fixture +def reset_stats(): + """Fixture to reset the global stats dict, mock disk writes, and reset throttle timestamps.""" + original_stats = router_main.stats + router_main.stats = copy.deepcopy(INITIAL_STATS) + + # Reset throttle timestamps to prevent state leakage between tests + original_last_stats_save = getattr(router_main, "_last_stats_save", 0.0) + original_last_save = getattr(router_main.record_tool_usage, "_last_save", 0.0) + router_main._last_stats_save = 0.0 + router_main.record_tool_usage._last_save = 0.0 + + with patch("router.main._atomic_write_json_sync") as mock_write: + yield mock_write + + router_main.stats = original_stats + router_main._last_stats_save = original_last_stats_save + router_main.record_tool_usage._last_save = original_last_save + +def test_record_tool_usage_basic_accumulation(reset_stats): + """Test that tool and global token counts are updated correctly.""" + router_main.record_tool_usage( + tool_name="shell", + prompt_tokens=10, + completion_tokens=20, + model="gpt-4", + latency_ms=100.5 + ) + + assert router_main.stats["tool_tokens"]["shell"] == 30 + assert router_main.stats["prompt_tokens"] == 10 + assert router_main.stats["completion_tokens"] == 20 + assert router_main.stats["routing_paths"]["litellm_fallback"] == 1 + + # Add another usage to test accumulation + router_main.record_tool_usage( + tool_name="shell", + prompt_tokens=5, + completion_tokens=15, + model="gpt-4", + latency_ms=50.0 + ) + + assert router_main.stats["tool_tokens"]["shell"] == 50 + assert router_main.stats["prompt_tokens"] == 15 + assert router_main.stats["completion_tokens"] == 35 + assert router_main.stats["routing_paths"]["litellm_fallback"] == 2 + +def test_record_tool_usage_none_fallback(reset_stats): + """Test that tool_name='none' is logged as 'other'.""" + router_main.record_tool_usage( + tool_name="none", + prompt_tokens=5, + completion_tokens=5, + model="gpt-4", + latency_ms=10.0 + ) + + assert router_main.stats["tool_tokens"]["other"] == 10 + assert "none" not in router_main.stats["tool_tokens"] + +def test_record_tool_usage_routing_paths_init(reset_stats): + """Test that routing_paths is initialized if missing and specific routes are tracked.""" + # Delete routing_paths to test initialization + del router_main.stats["routing_paths"] + + router_main.record_tool_usage( + tool_name="view", + prompt_tokens=0, + completion_tokens=0, + model="gpt-4", + latency_ms=0, + route="google_oauth_direct" + ) + + assert "routing_paths" in router_main.stats + assert router_main.stats["routing_paths"]["google_oauth_direct"] == 1 + assert router_main.stats["routing_paths"]["litellm_fallback"] == 0 + +@patch("router.main.time.strftime") +def test_record_tool_usage_timeline_buffer(mock_strftime, reset_stats): + """Test that events are appended and older events are dropped when limit > 15.""" + mock_strftime.return_value = "12:00:00" + + # Add 16 events + for i in range(16): + router_main.record_tool_usage( + tool_name=f"tool_{i}", + prompt_tokens=1, + completion_tokens=1, + model="test-model", + latency_ms=10.0 + i + ) + + timeline = router_main.stats["timeline"] + assert len(timeline) == 15 + + # The first event ("tool_0") should have been popped + assert timeline[0]["tool"] == "tool_1" + assert timeline[-1]["tool"] == "tool_15" + assert timeline[-1]["latency_ms"] == 25 + assert timeline[-1]["timestamp"] == "12:00:00" + +@patch("router.main.asyncio.get_running_loop") +def test_record_tool_usage_with_event_loop(mock_get_running_loop, reset_stats): + """Test background task creation when event loop is running.""" + mock_loop = MagicMock() + mock_get_running_loop.return_value = mock_loop + mock_task = MagicMock() + mock_loop.create_task.return_value = mock_task + + with patch("router.main._background_tasks", set()): + router_main.record_tool_usage("shell", 1, 1, "test", 10.0) + + mock_loop.create_task.assert_called_once() + mock_task.add_done_callback.assert_called_once() + +@patch("router.main.asyncio.get_running_loop") +def test_record_tool_usage_no_event_loop(mock_get_running_loop, reset_stats): + """Test fallback sync execution when no event loop is running.""" + mock_get_running_loop.side_effect = RuntimeError("no running event loop") + mock_atomic_write_sync = reset_stats + + router_main.record_tool_usage("shell", 1, 1, "test", 10.0) + + assert mock_atomic_write_sync.call_count == 2 # one for stats, one for timeline + +@patch("router.main.asyncio.get_running_loop") +def test_record_tool_usage_no_event_loop_fallback_error(mock_get_running_loop, reset_stats): + """Test fallback sync execution gracefully handles errors.""" + mock_get_running_loop.side_effect = RuntimeError("no running event loop") + mock_atomic_write_sync = reset_stats + mock_atomic_write_sync.side_effect = Exception("write failed") + + # Should not raise + router_main.record_tool_usage("shell", 1, 1, "test", 10.0) + + assert mock_atomic_write_sync.call_count == 2 From d75941055715c23dacaf1b08cbf267525c02bc60 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 24 Jun 2026 17:56:08 +0000 Subject: [PATCH 4/8] Add tests for memory_mcp.py Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com> From 79c4ac0ecf3d14994d7e0f4260761d60f387abdb Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 24 Jun 2026 18:54:11 +0000 Subject: [PATCH 5/8] Add tests for memory_mcp.py Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com> --- .github/workflows/test.yml | 2 +- .jules/bolt.md | 10 ++- pr_description.md | 13 --- pr_description.txt | 13 ++- router/main.py | 3 +- router/test_memory_mcp.py | 114 ------------------------ test_agy_proxy.py | 85 ------------------ test_circuit_breaker.py | 6 +- test_detect_active_tool.py | 114 ------------------------ test_memory_mcp.py | 0 test_record_tool_usage.py | 172 ------------------------------------- 11 files changed, 21 insertions(+), 511 deletions(-) delete mode 100644 pr_description.md delete mode 100644 router/test_memory_mcp.py delete mode 100644 test_agy_proxy.py delete mode 100644 test_detect_active_tool.py mode change 100755 => 100644 test_memory_mcp.py delete mode 100644 test_record_tool_usage.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index cd0f34de..3de6b85c 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -20,7 +20,7 @@ jobs: python-version: '3.11' - name: Install dependencies - run: pip install httpx==0.28.1 pytest pytest-asyncio redis + run: pip install httpx==0.28.1 - name: Run Circuit Breaker Tests run: python3 test_circuit_breaker.py diff --git a/.jules/bolt.md b/.jules/bolt.md index 9bda0f20..1a3987bd 100644 --- a/.jules/bolt.md +++ b/.jules/bolt.md @@ -1,3 +1,7 @@ -## 2024-06-24 - [Test implementation of detect_active_tool] -**Learning:** `router/main.py` requires configuration context (`CONFIG_PATH`) to be set, otherwise importing it for unit tests throws an error since it attempts to read config upon importing. -**Action:** Always mock or provide required environment variables such as `CONFIG_PATH` before importing functions from `router/main.py`. +## 2024-06-16 - Synchronous I/O in Async API Handlers +**Learning:** `save_persisted_stats()` was being called synchronously on every API request, cache hit, and tool usage log, triggering blocking disk I/O in the main event loop. +**Action:** Always throttle or batch background telemetry writes in async Python applications to prevent blocking the event loop under load. + +## 2026-06-24 - Async Offloading for SQLite I/O +**Learning:** Synchronous blocking I/O (like SQLite queries) inside an async event loop handler can cause significant latency spikes and block other requests. +**Action:** Use `asyncio.to_thread()` to offload synchronous database interactions to a worker thread, ensuring the main event loop remains responsive. diff --git a/pr_description.md b/pr_description.md deleted file mode 100644 index a15413ce..00000000 --- a/pr_description.md +++ /dev/null @@ -1,13 +0,0 @@ -# πŸ§ͺ Add tests for memory value JSON encoding/decoding - -## Description - -🎯 **What:** The testing gap addressed was an untested memory JSON encoder `_memory_value` and its counterpart decoder `_parse_memory_value` within `router/memory_mcp.py`. These pure functions handle payload storage preparation but lacked verification. - -πŸ“Š **Coverage:** The following scenarios are now tested: -- Happy paths for converting data and tags into serialized JSON strings. -- Edge cases where `tags` are provided as `None`. -- Unicode handling to ensure that `ensure_ascii=False` appropriately works without escaping characters. -- Decoding behavior of valid JSON, invalid JSON formats (catching `json.JSONDecodeError`), and TypeError handling for `_parse_memory_value`. - -✨ **Result:** Test coverage for the core memory data encoding mechanism is significantly improved, preventing potential regressions during memory data restructuring. Tests successfully passed `pytest`. diff --git a/pr_description.txt b/pr_description.txt index 26fc5f09..2f7c1980 100644 --- a/pr_description.txt +++ b/pr_description.txt @@ -1,3 +1,10 @@ -🎯 What: Add a comprehensive test file for `sync_gemini_token.py`. -πŸ“Š Coverage: Added tests covering the happy path (successful retrieval and parsing of credentials with nanosecond/offset handling), error conditions (secret-tool failure, empty output), edge cases for bad JSON (missing token key, missing access token), and exception handling (date parsing fallback, general exceptions). -✨ Result: `sync_gemini_token.py` is now fully covered by tests, ensuring its logic handles all expected scenarios correctly. +🎯 **What:** The `_memory_entry` helper function in `router/memory_mcp.py` was previously completely untested. This function is responsible for converting a raw dictionary representation of a memory entry from LiteLLM into a structured dictionary used by the MCP. + +πŸ“Š **Coverage:** The new `test_memory_mcp.py` file covers several scenarios for the `_memory_entry` dictionary translation helper: +- **Happy Path:** Tests that valid and complete memory dictionaries are properly parsed, mapped, and typed. +- **Invalid Key:** Verifies that if a memory has a key that does not start with `"memory:"`, the function properly returns `None`. +- **Malformed/String Value:** Tests the graceful fallback when a memory value is a raw string instead of the expected JSON payload, ensuring it still parses the string into the `data` field and initializes an empty `tags` array without throwing a JSON decode error. +- **Missing Fields:** Tests dictionary inputs that are missing either `"key"` or `"value"` or both to ensure graceful behavior, fallback values, and early exits without `KeyError`s. + +✨ **Result:** Increased code reliability and coverage for the core memory data transformation logic with pure, isolated dictionary tests ensuring zero runtime side effects. + diff --git a/router/main.py b/router/main.py index 69db4212..cd3de9d0 100644 --- a/router/main.py +++ b/router/main.py @@ -19,6 +19,7 @@ from pydantic import BaseModel from typing import Dict, Optional, Union + LITELLM_URL = (os.getenv("LITELLM_ADMIN_URL") or "http://127.0.0.1:4000").rstrip("/") LLAMA_SERVER_URL = (os.getenv("LLAMA_SERVER_URL") or "http://127.0.0.1:8080").rstrip("/") @@ -672,7 +673,7 @@ async def lifespan(app: FastAPI): get_http_client() await sync_cooldowns_from_valkey() - litellm_ready_url = f"{LITELLM_URL}/health/readiness" + litellm_ready_url = "http://127.0.0.1:4000/health/readiness" litellm_master_key = os.getenv("LITELLM_MASTER_KEY", "") max_wait = 180 logger.info(f"⏳ Waiting for LiteLLM on {LITELLM_URL} (max {max_wait}s)...") diff --git a/router/test_memory_mcp.py b/router/test_memory_mcp.py deleted file mode 100644 index 7bf535b8..00000000 --- a/router/test_memory_mcp.py +++ /dev/null @@ -1,114 +0,0 @@ -import pytest -import time -import re -import sys -import json -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent)) -from memory_mcp import _make_key, SCOPE_GLOBAL, SCOPE_LOCAL, PREFIX, _memory_value, _parse_memory_value - -def test_make_key_global(): - """Test generating a key for global scope.""" - category = "test_cat" - data = "test_data" - - before_ts = int(time.time() * 1000) - key = _make_key(category, True, data) - after_ts = int(time.time() * 1000) - - # Expected format: f"{PREFIX}:{scope}:{category}::{ts}:{h}" - parts = key.split(":") - - assert key.startswith(f"{PREFIX}:{SCOPE_GLOBAL}:{category}::") - - # Extract timestamp and hash part - # Format is memory:global:test_cat::1717612345:a1b2c3d4 - match = re.match(rf"^{PREFIX}:{SCOPE_GLOBAL}:{category}::(\d+):([a-z0-9x]+)$", key) - assert match is not None, f"Key {key} does not match expected format" - - ts = int(match.group(1)) - h = match.group(2) - - assert before_ts <= ts <= after_ts - assert len(h) <= 12 - -def test_make_key_local(): - """Test generating a key for local scope.""" - category = "another_cat" - data = "more_data" - - before_ts = int(time.time() * 1000) - key = _make_key(category, False, data) - after_ts = int(time.time() * 1000) - - assert key.startswith(f"{PREFIX}:{SCOPE_LOCAL}:{category}::") - - match = re.match(rf"^{PREFIX}:{SCOPE_LOCAL}:{category}::(\d+):([a-z0-9x]+)$", key) - assert match is not None, f"Key {key} does not match expected format" - - ts = int(match.group(1)) - h = match.group(2) - - assert before_ts <= ts <= after_ts - assert len(h) <= 12 - -def test_make_key_determinism_and_uniqueness(): - """Test determinism for same inputs within same timestamp, and uniqueness across timestamps/data.""" - category = "test_cat" - data1 = "data1" - data2 = "data2" - - key1 = _make_key(category, True, data1) - time.sleep(0.002) - key2 = _make_key(category, True, data1) - key3 = _make_key(category, True, data2) - - # Uniqueness across data - assert key1 != key3 - - # Check determinism: if the timestamp parts are the same, the keys should be identical - ts1 = key1.split("::")[1].split(":")[0] - ts2 = key2.split("::")[1].split(":")[0] - if ts1 == ts2: - assert key1 == key2 - else: - # If timestamp is different, keys should be different - assert key1 != key2 - - -def test_memory_value_happy_path(): - """Test _memory_value with standard data and tags.""" - result = _memory_value("some data", ["tag1", "tag2"]) - parsed = json.loads(result) - assert parsed == {"data": "some data", "tags": ["tag1", "tag2"]} - -def test_memory_value_missing_tags(): - """Test _memory_value when tags is None.""" - result = _memory_value("some data", None) - parsed = json.loads(result) - assert parsed == {"data": "some data", "tags": []} - -def test_memory_value_unicode(): - """Test _memory_value properly handles unicode and ensure_ascii=False.""" - result = _memory_value("こんにけは", ["δΈ–η•Œ"]) - # If ensure_ascii=False, the unicode characters shouldn't be escaped (no \uXXXX) - assert "こんにけは" in result - assert "δΈ–η•Œ" in result - parsed = json.loads(result) - assert parsed == {"data": "こんにけは", "tags": ["δΈ–η•Œ"]} - -def test_parse_memory_value_success(): - """Test _parse_memory_value successfully decodes valid JSON.""" - raw = '{"data": "info", "tags": ["a"]}' - result = _parse_memory_value(raw) - assert result == {"data": "info", "tags": ["a"]} - -def test_parse_memory_value_invalid_json(): - """Test _parse_memory_value with invalid JSON.""" - result = _parse_memory_value("{invalid_json:") - assert result == {"data": "{invalid_json:", "tags": []} - -def test_parse_memory_value_type_error(): - """Test _parse_memory_value with TypeError (e.g. passing None).""" - result = _parse_memory_value(None) - assert result == {"data": None, "tags": []} diff --git a/test_agy_proxy.py b/test_agy_proxy.py deleted file mode 100644 index 2e5219c6..00000000 --- a/test_agy_proxy.py +++ /dev/null @@ -1,85 +0,0 @@ -import pytest -from unittest.mock import patch, mock_open -import router.agy_proxy -from router.agy_proxy import _is_quota_exhausted - -@pytest.fixture(autouse=True) -def reset_global_throttle(): - """Reset the global throttle variable before each test.""" - router.agy_proxy._last_log_check = 0.0 - -def test_direct_stderr_exhaustion(): - assert _is_quota_exhausted(1, "", "RESOURCE_EXHAUSTED: out of quota") - assert _is_quota_exhausted(1, "", "error code 429 received") - assert _is_quota_exhausted(1, "", "your quota reached the limit") - assert _is_quota_exhausted(1, "", "rate limit exceeded") - -def test_success_response(): - assert not _is_quota_exhausted(0, "success response", "") - assert not _is_quota_exhausted(0, "success", "some warning") - -def test_non_quota_error(): - assert not _is_quota_exhausted(1, "", "syntax error") - assert not _is_quota_exhausted(127, "", "command not found") - -@patch("os.path.exists") -@patch("time.time") -def test_silent_failure_fallback(mock_time, mock_exists): - # Empty stdout/stderr with rc=0 but no log file - mock_exists.return_value = False - mock_time.return_value = 100.0 - - assert _is_quota_exhausted(0, "", "") == True - -@patch("router.agy_proxy.open", new_callable=mock_open, read_data="some log\nRESOURCE_EXHAUSTED\n") -@patch("os.path.exists") -@patch("time.time") -def test_silent_failure_with_log_exhaustion(mock_time, mock_exists, mock_file): - mock_exists.return_value = True - mock_time.return_value = 100.0 - - assert _is_quota_exhausted(0, "", "") == True - mock_file.assert_called_once() - -@patch("router.agy_proxy.open", new_callable=mock_open, read_data="normal log\nno errors\n") -@patch("os.path.exists") -@patch("time.time") -def test_silent_failure_with_log_no_exhaustion(mock_time, mock_exists, mock_file): - mock_exists.return_value = True - mock_time.return_value = 100.0 - - # Returns True even if log doesn't have it because "Empty stdout+stderr with rc=0 strongly suggests quota exhaustion" - assert _is_quota_exhausted(0, "", "") == True - mock_file.assert_called_once() - -@patch("router.agy_proxy.open", new_callable=mock_open, read_data="RESOURCE_EXHAUSTED\n") -@patch("os.path.exists") -@patch("time.time") -def test_throttling_behavior(mock_time, mock_exists, mock_file): - mock_exists.return_value = True - mock_time.return_value = 100.0 - - # First call - should read file - assert _is_quota_exhausted(0, "", "") == True - assert mock_file.call_count == 1 - - # Second call immediately after (time = 101.0, diff = 1s < 2s) - mock_time.return_value = 101.0 - assert _is_quota_exhausted(0, "", "") == True - assert mock_file.call_count == 1 # File not read again due to throttle - - # Third call after 2 seconds (time = 103.0, diff = 3s > 2s) - mock_time.return_value = 103.0 - assert _is_quota_exhausted(0, "", "") == True - assert mock_file.call_count == 2 # File read again - -@patch("router.agy_proxy.open") -@patch("os.path.exists") -@patch("time.time") -def test_log_read_exception(mock_time, mock_exists, mock_file): - mock_exists.return_value = True - mock_time.return_value = 100.0 - mock_file.side_effect = Exception("Permission denied") - - # Should catch exception and still return True as fallback - assert _is_quota_exhausted(0, "", "") == True diff --git a/test_circuit_breaker.py b/test_circuit_breaker.py index 2c44defa..16856232 100644 --- a/test_circuit_breaker.py +++ b/test_circuit_breaker.py @@ -139,7 +139,6 @@ def test_backward_compatibility(): print("βœ“ Master record_failure and record_success maintain compatibility") - def test_dual_breaker_tier_max_logic(): """Master breaker tier returns max of sub-breakers.""" reset_breakers() @@ -198,7 +197,6 @@ def test_full_cycle(): print("βœ“ Full cycle: 3 failures β†’ Tier 3 β†’ probe success β†’ reset") - def test_sync_from_valkey_exception_handling(): """Exception during Valkey sync is caught and logged.""" reset_breakers() @@ -270,8 +268,6 @@ async def test_save_to_valkey_exception_handling(): with patch('router.circuit_breaker.logger') as mock_logger: await sub.save_to_valkey(mock_redis) mock_logger.warning.assert_called_once() - - if __name__ == "__main__": test_initial_state() test_first_failure_trips_to_tier1() @@ -287,7 +283,7 @@ async def test_save_to_valkey_exception_handling(): asyncio.run(test_save_to_valkey_success()) asyncio.run(test_save_to_valkey_no_client()) asyncio.run(test_save_to_valkey_exception_handling()) - + print("\n" + "=" * 60) print(" ALL CIRCUIT BREAKER TESTS PASSED βœ“") print("=" * 60) diff --git a/test_detect_active_tool.py b/test_detect_active_tool.py deleted file mode 100644 index c06edbe4..00000000 --- a/test_detect_active_tool.py +++ /dev/null @@ -1,114 +0,0 @@ -import os - -# Resolve the absolute path to the config file to ensure tests can run from any working directory -base_dir = os.path.dirname(os.path.abspath(__file__)) -os.environ["CONFIG_PATH"] = os.path.join(base_dir, "router", "config.yaml") - -import pytest -from router.main import detect_active_tool, map_tool_to_category - -@pytest.mark.parametrize( - "tool_name,expected_category", - [ - ("tree", "tree"), - ("list_dir", "tree"), - ("list-dir", "tree"), - ("shell", "shell"), - ("command", "shell"), - ("run", "shell"), - ("execute", "shell"), - ("cmd", "shell"), - ("write", "write"), - ("create", "write"), - ("save", "write"), - ("edit", "write"), - ("patch", "write"), - ("replace", "write"), - ("view", "view"), - ("cat", "view"), - ("grep", "view"), - ("read", "view"), - ("search", "view"), - ("find", "view"), - ("something_else", "other"), - ("unknown__list_dir", "tree"), - ], -) -def test_map_tool_to_category(tool_name, expected_category): - """Verify tool names are mapped correctly to high-level categories.""" - assert map_tool_to_category(tool_name) == expected_category - -def test_detect_active_tool_tool_role_with_name(): - """Verify detect_active_tool correctly identifies a tool when role is 'tool' and name is provided.""" - payload = { - "messages": [ - {"role": "tool", "name": "run_command", "content": "file1.txt\nfile2.txt"} - ] - } - assert detect_active_tool(payload) == "shell" - -def test_detect_active_tool_tool_role_without_name(): - """Verify detect_active_tool looks backward for the assistant tool request when name is absent.""" - payload = { - "messages": [ - {"role": "user", "content": "run this cmd"}, - {"role": "assistant", "tool_calls": [ - {"id": "call_1", "function": {"name": "run_command"}} - ]}, - {"role": "tool", "tool_call_id": "call_1", "content": "output"} - ] - } - assert detect_active_tool(payload) == "shell" - -def test_detect_active_tool_tool_role_without_name_no_match(): - """Verify detect_active_tool handles the case where it can't find the backward assistant tool request.""" - payload = { - "messages": [ - {"role": "user", "content": "run this cmd"}, - {"role": "assistant", "tool_calls": [ - {"id": "call_2", "function": {"name": "run_command"}} - ]}, - {"role": "tool", "tool_call_id": "call_1", "content": "output"} - ] - } - assert detect_active_tool(payload) == "other" - -def test_detect_active_tool_assistant_role(): - """Verify detect_active_tool correctly identifies a tool when role is 'assistant' with tool_calls.""" - payload = { - "messages": [ - {"role": "assistant", "tool_calls": [ - {"id": "call_1", "function": {"name": "write_file"}} - ]} - ] - } - assert detect_active_tool(payload) == "write" - -@pytest.mark.parametrize( - "user_message,expected_category", - [ - ("tree", "tree"), - ("files", "tree"), - ("shell", "shell"), - ("run", "shell"), - ("cmd", "shell"), - ("write", "write"), - ("create file", "write"), - ("view", "view"), - ("read", "view"), - ("cat", "view"), - ], -) -def test_detect_active_tool_user_fallback(user_message, expected_category): - """Verify detect_active_tool falls back to keyword matching in user messages.""" - assert detect_active_tool({"messages": [{"role": "user", "content": user_message}]}) == expected_category - -def test_detect_active_tool_empty_and_malformed(): - """Verify detect_active_tool handles empty or malformed inputs correctly.""" - assert detect_active_tool({}) == "none" - assert detect_active_tool({"messages": []}) == "none" - assert detect_active_tool({"messages": ["not a dict"]}) == "none" - assert detect_active_tool({"messages": [{"role": "user", "content": "hello"}]}) == "none" - # test nested malformed objects - assert detect_active_tool({"messages": [{"role": "assistant", "tool_calls": ["not a dict"]}]}) == "none" - assert detect_active_tool({"messages": [{"role": "assistant", "tool_calls": [{"id": "call_1", "function": "not a dict"}]}]}) == "other" diff --git a/test_memory_mcp.py b/test_memory_mcp.py old mode 100755 new mode 100644 diff --git a/test_record_tool_usage.py b/test_record_tool_usage.py deleted file mode 100644 index db9ab362..00000000 --- a/test_record_tool_usage.py +++ /dev/null @@ -1,172 +0,0 @@ -import pytest -import asyncio -from unittest.mock import patch, MagicMock, call -import copy - -import router.main as router_main - -# A snapshot of the initial stats structure for resetting -INITIAL_STATS = { - "total_requests": 0, - "simple_requests": 0, - "medium_requests": 0, - "complex_requests": 0, - "reasoning_requests": 0, - "advanced_requests": 0, - "cache_hits": 0, - "last_triage_decision": "None", - "avg_triage_latency_ms": 0.0, - "avg_proxy_latency_ms": 0.0, - "total_triage_time_ms": 0.0, - "total_proxy_time_ms": 0.0, - "prompt_tokens": 0, - "completion_tokens": 0, - "tool_tokens": { - "tree": 0, - "shell": 0, - "write": 0, - "view": 0, - "other": 0 - }, - "routing_paths": {"google_oauth_direct": 0, "litellm_fallback": 0}, - "timeline": [] -} - -@pytest.fixture -def reset_stats(): - """Fixture to reset the global stats dict, mock disk writes, and reset throttle timestamps.""" - original_stats = router_main.stats - router_main.stats = copy.deepcopy(INITIAL_STATS) - - # Reset throttle timestamps to prevent state leakage between tests - original_last_stats_save = getattr(router_main, "_last_stats_save", 0.0) - original_last_save = getattr(router_main.record_tool_usage, "_last_save", 0.0) - router_main._last_stats_save = 0.0 - router_main.record_tool_usage._last_save = 0.0 - - with patch("router.main._atomic_write_json_sync") as mock_write: - yield mock_write - - router_main.stats = original_stats - router_main._last_stats_save = original_last_stats_save - router_main.record_tool_usage._last_save = original_last_save - -def test_record_tool_usage_basic_accumulation(reset_stats): - """Test that tool and global token counts are updated correctly.""" - router_main.record_tool_usage( - tool_name="shell", - prompt_tokens=10, - completion_tokens=20, - model="gpt-4", - latency_ms=100.5 - ) - - assert router_main.stats["tool_tokens"]["shell"] == 30 - assert router_main.stats["prompt_tokens"] == 10 - assert router_main.stats["completion_tokens"] == 20 - assert router_main.stats["routing_paths"]["litellm_fallback"] == 1 - - # Add another usage to test accumulation - router_main.record_tool_usage( - tool_name="shell", - prompt_tokens=5, - completion_tokens=15, - model="gpt-4", - latency_ms=50.0 - ) - - assert router_main.stats["tool_tokens"]["shell"] == 50 - assert router_main.stats["prompt_tokens"] == 15 - assert router_main.stats["completion_tokens"] == 35 - assert router_main.stats["routing_paths"]["litellm_fallback"] == 2 - -def test_record_tool_usage_none_fallback(reset_stats): - """Test that tool_name='none' is logged as 'other'.""" - router_main.record_tool_usage( - tool_name="none", - prompt_tokens=5, - completion_tokens=5, - model="gpt-4", - latency_ms=10.0 - ) - - assert router_main.stats["tool_tokens"]["other"] == 10 - assert "none" not in router_main.stats["tool_tokens"] - -def test_record_tool_usage_routing_paths_init(reset_stats): - """Test that routing_paths is initialized if missing and specific routes are tracked.""" - # Delete routing_paths to test initialization - del router_main.stats["routing_paths"] - - router_main.record_tool_usage( - tool_name="view", - prompt_tokens=0, - completion_tokens=0, - model="gpt-4", - latency_ms=0, - route="google_oauth_direct" - ) - - assert "routing_paths" in router_main.stats - assert router_main.stats["routing_paths"]["google_oauth_direct"] == 1 - assert router_main.stats["routing_paths"]["litellm_fallback"] == 0 - -@patch("router.main.time.strftime") -def test_record_tool_usage_timeline_buffer(mock_strftime, reset_stats): - """Test that events are appended and older events are dropped when limit > 15.""" - mock_strftime.return_value = "12:00:00" - - # Add 16 events - for i in range(16): - router_main.record_tool_usage( - tool_name=f"tool_{i}", - prompt_tokens=1, - completion_tokens=1, - model="test-model", - latency_ms=10.0 + i - ) - - timeline = router_main.stats["timeline"] - assert len(timeline) == 15 - - # The first event ("tool_0") should have been popped - assert timeline[0]["tool"] == "tool_1" - assert timeline[-1]["tool"] == "tool_15" - assert timeline[-1]["latency_ms"] == 25 - assert timeline[-1]["timestamp"] == "12:00:00" - -@patch("router.main.asyncio.get_running_loop") -def test_record_tool_usage_with_event_loop(mock_get_running_loop, reset_stats): - """Test background task creation when event loop is running.""" - mock_loop = MagicMock() - mock_get_running_loop.return_value = mock_loop - mock_task = MagicMock() - mock_loop.create_task.return_value = mock_task - - with patch("router.main._background_tasks", set()): - router_main.record_tool_usage("shell", 1, 1, "test", 10.0) - - mock_loop.create_task.assert_called_once() - mock_task.add_done_callback.assert_called_once() - -@patch("router.main.asyncio.get_running_loop") -def test_record_tool_usage_no_event_loop(mock_get_running_loop, reset_stats): - """Test fallback sync execution when no event loop is running.""" - mock_get_running_loop.side_effect = RuntimeError("no running event loop") - mock_atomic_write_sync = reset_stats - - router_main.record_tool_usage("shell", 1, 1, "test", 10.0) - - assert mock_atomic_write_sync.call_count == 2 # one for stats, one for timeline - -@patch("router.main.asyncio.get_running_loop") -def test_record_tool_usage_no_event_loop_fallback_error(mock_get_running_loop, reset_stats): - """Test fallback sync execution gracefully handles errors.""" - mock_get_running_loop.side_effect = RuntimeError("no running event loop") - mock_atomic_write_sync = reset_stats - mock_atomic_write_sync.side_effect = Exception("write failed") - - # Should not raise - router_main.record_tool_usage("shell", 1, 1, "test", 10.0) - - assert mock_atomic_write_sync.call_count == 2 From 81b8255194204fc789dc9807c95e6cf2b211679a Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 24 Jun 2026 20:03:53 +0000 Subject: [PATCH 6/8] Add tests for memory_mcp.py Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com> --- router/main.py | 6 +- router/test_memory_mcp.py | 113 ++++++++++++++++++++++++++++++++++++++ test_atomic_write.py | 1 + test_circuit_breaker.py | 7 ++- 4 files changed, 120 insertions(+), 7 deletions(-) create mode 100644 router/test_memory_mcp.py diff --git a/router/main.py b/router/main.py index cd3de9d0..3c149b0f 100644 --- a/router/main.py +++ b/router/main.py @@ -18,8 +18,6 @@ from circuit_breaker import get_breaker from pydantic import BaseModel from typing import Dict, Optional, Union - - LITELLM_URL = (os.getenv("LITELLM_ADMIN_URL") or "http://127.0.0.1:4000").rstrip("/") LLAMA_SERVER_URL = (os.getenv("LLAMA_SERVER_URL") or "http://127.0.0.1:8080").rstrip("/") @@ -673,7 +671,7 @@ async def lifespan(app: FastAPI): get_http_client() await sync_cooldowns_from_valkey() - litellm_ready_url = "http://127.0.0.1:4000/health/readiness" + litellm_ready_url = f"{LITELLM_URL}/health/readiness" litellm_master_key = os.getenv("LITELLM_MASTER_KEY", "") max_wait = 180 logger.info(f"⏳ Waiting for LiteLLM on {LITELLM_URL} (max {max_wait}s)...") @@ -1367,7 +1365,7 @@ async def proxy_memory(request: Request, path: str = ""): # Exclude standard headers that FastAPI/uvicorn will manage for h in ["content-encoding", "content-length", "transfer-encoding", "connection"]: response_headers.pop(h, None) - + return Response( content=r.content, status_code=r.status_code, diff --git a/router/test_memory_mcp.py b/router/test_memory_mcp.py new file mode 100644 index 00000000..49dc3355 --- /dev/null +++ b/router/test_memory_mcp.py @@ -0,0 +1,113 @@ +import pytest +import time +import re +import sys +import json +from pathlib import Path +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from memory_mcp import _make_key, SCOPE_GLOBAL, SCOPE_LOCAL, PREFIX, _memory_value, _parse_memory_value + +def test_make_key_global(): + """Test generating a key for global scope.""" + category = "test_cat" + data = "test_data" + + before_ts = int(time.time() * 1000) + key = _make_key(category, True, data) + after_ts = int(time.time() * 1000) + + # Expected format: f"{PREFIX}:{scope}:{category}::{ts}:{h}" + parts = key.split(":") + + assert key.startswith(f"{PREFIX}:{SCOPE_GLOBAL}:{category}::") + + # Extract timestamp and hash part + # Format is memory:global:test_cat::1717612345:a1b2c3d4 + match = re.match(rf"^{PREFIX}:{SCOPE_GLOBAL}:{category}::(\d+):([a-z0-9x]+)$", key) + assert match is not None, f"Key {key} does not match expected format" + + ts = int(match.group(1)) + h = match.group(2) + + assert before_ts <= ts <= after_ts + assert len(h) <= 12 + +def test_make_key_local(): + """Test generating a key for local scope.""" + category = "another_cat" + data = "more_data" + + before_ts = int(time.time() * 1000) + key = _make_key(category, False, data) + after_ts = int(time.time() * 1000) + + assert key.startswith(f"{PREFIX}:{SCOPE_LOCAL}:{category}::") + + match = re.match(rf"^{PREFIX}:{SCOPE_LOCAL}:{category}::(\d+):([a-z0-9x]+)$", key) + assert match is not None, f"Key {key} does not match expected format" + + ts = int(match.group(1)) + h = match.group(2) + + assert before_ts <= ts <= after_ts + assert len(h) <= 12 + +def test_make_key_determinism_and_uniqueness(): + """Test determinism for same inputs within same timestamp, and uniqueness across timestamps/data.""" + category = "test_cat" + data1 = "data1" + data2 = "data2" + + key1 = _make_key(category, True, data1) + time.sleep(0.002) + key2 = _make_key(category, True, data1) + key3 = _make_key(category, True, data2) + + # Uniqueness across data + assert key1 != key3 + + # Check determinism: if the timestamp parts are the same, the keys should be identical + ts1 = key1.split("::")[1].split(":")[0] + ts2 = key2.split("::")[1].split(":")[0] + if ts1 == ts2: + assert key1 == key2 + else: + # If timestamp is different, keys should be different + assert key1 != key2 + +def test_memory_value_happy_path(): + """Test _memory_value with standard data and tags.""" + result = _memory_value("some data", ["tag1", "tag2"]) + parsed = json.loads(result) + assert parsed == {"data": "some data", "tags": ["tag1", "tag2"]} + +def test_memory_value_missing_tags(): + """Test _memory_value when tags is None.""" + result = _memory_value("some data", None) + parsed = json.loads(result) + assert parsed == {"data": "some data", "tags": []} + +def test_memory_value_unicode(): + """Test _memory_value properly handles unicode and ensure_ascii=False.""" + result = _memory_value("こんにけは", ["δΈ–η•Œ"]) + # If ensure_ascii=False, the unicode characters shouldn't be escaped (no \uXXXX) + assert "こんにけは" in result + assert "δΈ–η•Œ" in result + parsed = json.loads(result) + assert parsed == {"data": "こんにけは", "tags": ["δΈ–η•Œ"]} + +def test_parse_memory_value_success(): + """Test _parse_memory_value successfully decodes valid JSON.""" + raw = '{"data": "info", "tags": ["a"]}' + result = _parse_memory_value(raw) + assert result == {"data": "info", "tags": ["a"]} + +def test_parse_memory_value_invalid_json(): + """Test _parse_memory_value with invalid JSON.""" + result = _parse_memory_value("{invalid_json:") + assert result == {"data": "{invalid_json:", "tags": []} + +def test_parse_memory_value_type_error(): + """Test _parse_memory_value with TypeError (e.g. passing None).""" + result = _parse_memory_value(None) + assert result == {"data": None, "tags": []} diff --git a/test_atomic_write.py b/test_atomic_write.py index 2d64163e..1d9b7cc6 100644 --- a/test_atomic_write.py +++ b/test_atomic_write.py @@ -210,3 +210,4 @@ class Unserializable: loaded_data = json.load(f) assert loaded_data == old_data + diff --git a/test_circuit_breaker.py b/test_circuit_breaker.py index 16856232..e1635f42 100644 --- a/test_circuit_breaker.py +++ b/test_circuit_breaker.py @@ -202,6 +202,7 @@ def test_sync_from_valkey_exception_handling(): reset_breakers() b = get_breaker() + mock_redis = AsyncMock() mock_redis.hgetall.side_effect = Exception("Simulated connection error") @@ -215,7 +216,7 @@ def test_sync_from_valkey_exception_handling(): print("βœ“ Valkey sync exception handling") -@pytest.mark.asyncio +@pytest.mark.anyio async def test_save_to_valkey_success(): """Verify state is correctly serialized and persisted to Valkey.""" b = get_breaker() @@ -246,7 +247,7 @@ async def test_save_to_valkey_success(): print("βœ“ Valkey save succeeds with correct data and TTL") -@pytest.mark.asyncio +@pytest.mark.anyio async def test_save_to_valkey_no_client(): """Verify early return when redis client is None.""" b = get_breaker() @@ -256,7 +257,7 @@ async def test_save_to_valkey_no_client(): print("βœ“ Valkey save handles None client safely") -@pytest.mark.asyncio +@pytest.mark.anyio async def test_save_to_valkey_exception_handling(): """Verify exceptions during Valkey save are caught and logged.""" b = get_breaker() From 79216d541534131bd044a8220ede039a00edad05 Mon Sep 17 00:00:00 2001 From: sheepdestroyer Date: Wed, 24 Jun 2026 22:24:50 +0200 Subject: [PATCH 7/8] Resolve merge conflicts, integrate pytest to CI, and address review feedback --- .github/workflows/test.yml | 7 ++++--- pytest.ini | 4 ++++ router/agy_proxy.py | 6 ++++-- router/main.py | 18 ++++++++++++------ test_models_proxy.py | 31 +++++++++++++++++++++++++------ test_sync_gemini_token.py | 7 ++++++- 6 files changed, 55 insertions(+), 18 deletions(-) create mode 100644 pytest.ini diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 3de6b85c..ac8d6823 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -20,13 +20,14 @@ jobs: python-version: '3.11' - name: Install dependencies - run: pip install httpx==0.28.1 + run: pip install httpx==0.28.1 pytest anyio pyyaml - - name: Run Circuit Breaker Tests - run: python3 test_circuit_breaker.py + - name: Run Unit Tests + run: CONFIG_PATH=router/config.yaml PYTHONPATH=.:router pytest --ignore=test_agy_behavior.py --ignore=test_agy_tiers.py - name: Run Breaker Verification run: python3 verify_breaker.py - name: Run Integration Verification run: python3 test_a2_verify.py + diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 00000000..a0525e89 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,4 @@ +[pytest] +addopts = --import-mode=importlib +norecursedirs = clickhouse-data postgres-data langfuse-data redis-lf-data valkey-data minio-data .git .github + diff --git a/router/agy_proxy.py b/router/agy_proxy.py index 16adb8e9..5788e732 100644 --- a/router/agy_proxy.py +++ b/router/agy_proxy.py @@ -69,6 +69,8 @@ async def save(self) -> None: # agy_conversation_data = {"conversation_id": str, "current_tier_index": int} _session_store: dict = {} +AGY_DAEMON_URL = os.environ.get("AGY_DAEMON_URL", "http://127.0.0.1:5005") + def _get_last_conversation_id() -> Optional[str]: """Read the last conversation ID for our workspace from agy's cache file.""" @@ -91,7 +93,7 @@ async def _run_agy_print(client: httpx.AsyncClient, prompt: str, model_override: """ Forward the agy execution request to the host-side agy daemon. """ - url = "http://127.0.0.1:5005/run" + url = f"{AGY_DAEMON_URL}/run" payload = { "prompt": prompt, "model_override": model_override, @@ -299,7 +301,7 @@ async def try_agy_proxy(prompt: str, messages: list = None, tier_timeout = min(AGY_TIMEOUT_SECS, remaining) if stream: - url = "http://127.0.0.1:5005/run" + url = f"{AGY_DAEMON_URL}/run" payload = { "prompt": proxy_prompt, "model_override": tier["env_override"], diff --git a/router/main.py b/router/main.py index 3c149b0f..2d43966a 100644 --- a/router/main.py +++ b/router/main.py @@ -400,7 +400,7 @@ async def sync_adaptive_router_roster(master_key: str): logger.warning("No LITELLM_MASTER_KEY β€” skipping roster sync") return headers = {"Authorization": f"Bearer {master_key}", "Content-Type": "application/json"} - admin_url = "http://127.0.0.1:4000" + admin_url = LITELLM_URL try: client = get_http_client() r = await client.get("https://openrouter.ai/api/v1/models", timeout=5.0) @@ -558,7 +558,7 @@ async def _register_ollama_models_in_db(master_key: str): logger.warning("No LiteLLM master key provided β€” skipping Ollama DB registration") return - admin_url = os.getenv("LITELLM_ADMIN_URL", "http://127.0.0.1:4000") + admin_url = LITELLM_URL headers = {"Authorization": f"Bearer {master_key}", "Content-Type": "application/json"} ollama_models = [] @@ -1151,8 +1151,14 @@ async def get_llamacpp_metrics() -> dict: if r3.status_code == 200: slots_data = r3.json() for s in slots_data: - next_tok = s.get("next_token", [{}]) - decoded = next_tok[0].get("n_decoded", 0) if next_tok else 0 + next_tok = s.get("next_token") + decoded = 0 + if isinstance(next_tok, dict): + decoded = next_tok.get("n_decoded", 0) + elif isinstance(next_tok, list) and next_tok: + first_tok = next_tok[0] + if isinstance(first_tok, dict): + decoded = first_tok.get("n_decoded", 0) result["slots"].append({ "id": s.get("id", 0), "is_processing": s.get("is_processing", False), @@ -1206,7 +1212,7 @@ def _save_free_models_roster(free_models: list[dict]) -> None: import datetime as _dt payload = { "models": free_models, - "updated_at": _dt.datetime.utcnow().isoformat() + "Z", + "updated_at": _dt.datetime.now(_dt.timezone.utc).isoformat().replace("+00:00", "Z"), "count": len(free_models) } try: @@ -1220,7 +1226,7 @@ def _save_best_model_to_disk(best_model: dict) -> None: """Persist the best free model to a JSON file Ralph can read.""" import json as _json import datetime as _dt - payload = {**best_model, "updated_at": _dt.datetime.utcnow().isoformat() + "Z"} + payload = {**best_model, "updated_at": _dt.datetime.now(_dt.timezone.utc).isoformat().replace("+00:00", "Z")} try: with open("/config/router_dir/best_free_model.json", "w") as f: _json.dump(payload, f, indent=2) diff --git a/test_models_proxy.py b/test_models_proxy.py index f67743f1..af95a188 100644 --- a/test_models_proxy.py +++ b/test_models_proxy.py @@ -14,12 +14,31 @@ from main import get_http_client, proxy_models, HTTP_MAX_CONNECTIONS, HTTP_MAX_KEEPALIVE_CONNECTIONS, HTTP_KEEPALIVE_EXPIRY def test_http_client_limits(): - # Verify that get_http_client initializes with configured limits - client = get_http_client() - pool = client._transport._pool - assert pool._max_connections == HTTP_MAX_CONNECTIONS - assert pool._max_keepalive_connections == HTTP_MAX_KEEPALIVE_CONNECTIONS - assert pool._keepalive_expiry == HTTP_KEEPALIVE_EXPIRY + # Verify that get_http_client initializes with configured limits using public mocks + import main + import httpx + + original_init = httpx.Limits.__init__ + calls = [] + + def spy_init(self, *args, **kwargs): + calls.append((args, kwargs)) + original_init(self, *args, **kwargs) + + original_client = main._http_client + main._http_client = None + try: + with patch.object(httpx.Limits, "__init__", new=spy_init): + main.get_http_client() + assert len(calls) == 1 + args, kwargs = calls[0] + assert kwargs.get("max_connections") == main.HTTP_MAX_CONNECTIONS + assert kwargs.get("max_keepalive_connections") == main.HTTP_MAX_KEEPALIVE_CONNECTIONS + assert kwargs.get("keepalive_expiry") == main.HTTP_KEEPALIVE_EXPIRY + finally: + main._http_client = original_client + + @pytest.mark.anyio async def test_proxy_models_success(): diff --git a/test_sync_gemini_token.py b/test_sync_gemini_token.py index ada9f95a..d0cbe351 100644 --- a/test_sync_gemini_token.py +++ b/test_sync_gemini_token.py @@ -1,4 +1,5 @@ import json +import os import pytest from unittest.mock import patch, mock_open, MagicMock import sys @@ -49,7 +50,7 @@ def test_happy_path(mock_subprocess, mock_os_makedirs, mock_time, capsys): text=True ) - mock_os_makedirs.assert_called_once() + mock_os_makedirs.assert_called_once_with(os.path.dirname(sync_gemini_token.TARGET_PATH), exist_ok=True) m_open.assert_called_once_with(sync_gemini_token.TARGET_PATH, "w") # Check the written file @@ -136,6 +137,8 @@ def test_fallback_expiry(mock_subprocess, mock_os_makedirs, mock_time, capsys): with patch('builtins.open', m_open): sync_gemini_token.main() + mock_os_makedirs.assert_called_once_with(os.path.dirname(sync_gemini_token.TARGET_PATH), exist_ok=True) + # Assert standard output/error messages using capsys captured = capsys.readouterr() assert "Warning: Failed to parse expiry date 'invalid-date'" in captured.err @@ -170,6 +173,8 @@ def test_expired_token(mock_subprocess, mock_os_makedirs, mock_time, capsys): with patch('builtins.open', m_open): sync_gemini_token.main() + mock_os_makedirs.assert_called_once_with(os.path.dirname(sync_gemini_token.TARGET_PATH), exist_ok=True) + # Assert standard output message using capsys for expired token captured = capsys.readouterr() # 1600 seconds = 26m 40s ago From 9df83e471cf4a854cbc024fdfeda9f701cd02ebb Mon Sep 17 00:00:00 2001 From: sheepdestroyer Date: Wed, 24 Jun 2026 22:26:09 +0200 Subject: [PATCH 8/8] Install all required project dependencies in CI workflow --- .github/workflows/test.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index ac8d6823..bb00917a 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -20,7 +20,7 @@ jobs: python-version: '3.11' - name: Install dependencies - run: pip install httpx==0.28.1 pytest anyio pyyaml + run: pip install httpx==0.28.1 pytest anyio pyyaml fastapi uvicorn python-multipart asyncpg langfuse redis - name: Run Unit Tests run: CONFIG_PATH=router/config.yaml PYTHONPATH=.:router pytest --ignore=test_agy_behavior.py --ignore=test_agy_tiers.py