diff --git a/tests/unit/test_concurrency/__init__.py b/tests/unit/test_concurrency/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/tests/unit/test_concurrency/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/unit/test_concurrency/test_consumer_concurrency.py b/tests/unit/test_concurrency/test_consumer_concurrency.py new file mode 100644 index 00000000..32a6559b --- /dev/null +++ b/tests/unit/test_concurrency/test_consumer_concurrency.py @@ -0,0 +1,286 @@ +""" +Tests for Consumer concurrency: TaskGroup-based concurrent message processing, +rate-limit retry with backpressure, and message acknowledgement. +""" + +import asyncio +import time + +import pytest +from unittest.mock import MagicMock, AsyncMock, patch + +from trustgraph.base.consumer import Consumer +from trustgraph.exceptions import TooManyRequests + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_consumer( + concurrency=1, + handler=None, + rate_limit_retry_time=0.01, + rate_limit_timeout=1, +): + """Create a Consumer with mocked infrastructure.""" + taskgroup = MagicMock() + flow = MagicMock() + backend = MagicMock() + schema = MagicMock() + handler = handler or AsyncMock() + + consumer = Consumer( + taskgroup=taskgroup, + flow=flow, + backend=backend, + topic="test-topic", + subscriber="test-sub", + schema=schema, + handler=handler, + rate_limit_retry_time=rate_limit_retry_time, + rate_limit_timeout=rate_limit_timeout, + concurrency=concurrency, + ) + + return consumer + + +def _make_msg(): + """Create a mock Pulsar message.""" + return MagicMock() + + +# --------------------------------------------------------------------------- +# Concurrency configuration tests +# --------------------------------------------------------------------------- + +class TestConcurrencyConfiguration: + + def test_default_concurrency_is_1(self): + consumer = _make_consumer() + assert consumer.concurrency == 1 + + def test_custom_concurrency(self): + consumer = _make_consumer(concurrency=10) + assert consumer.concurrency == 10 + + def test_concurrency_stored(self): + for n in [1, 5, 20, 100]: + consumer = _make_consumer(concurrency=n) + assert consumer.concurrency == n + + +class TestTaskGroupConcurrency: + + @pytest.mark.asyncio + async def test_creates_n_concurrent_tasks(self): + """consumer_run should create exactly N concurrent consume_from_queue tasks.""" + concurrency = 5 + consumer = _make_consumer(concurrency=concurrency) + + # Track how many consume_from_queue calls are made + call_count = 0 + original_running = True + + async def mock_consume(): + nonlocal call_count + call_count += 1 + # Wait a bit to let all tasks start, then signal stop + await asyncio.sleep(0.05) + consumer.running = False + + consumer.consume_from_queue = mock_consume + + # Mock the backend.create_consumer + consumer.backend.create_consumer = MagicMock(return_value=MagicMock()) + + # Run consumer_run - it will create TaskGroup with N tasks + consumer.running = True + await consumer.consumer_run() + + assert call_count == concurrency + + @pytest.mark.asyncio + async def test_single_concurrency_creates_one_task(self): + """With concurrency=1, only one consume_from_queue task is created.""" + consumer = _make_consumer(concurrency=1) + call_count = 0 + + async def mock_consume(): + nonlocal call_count + call_count += 1 + await asyncio.sleep(0.01) + consumer.running = False + + consumer.consume_from_queue = mock_consume + consumer.backend.create_consumer = MagicMock(return_value=MagicMock()) + + consumer.running = True + await consumer.consumer_run() + + assert call_count == 1 + + +# --------------------------------------------------------------------------- +# Rate-limit retry tests +# --------------------------------------------------------------------------- + +class TestRateLimitRetry: + + @pytest.mark.asyncio + async def test_rate_limit_retries_then_succeeds(self): + """TooManyRequests should cause retry, then succeed on next attempt.""" + call_count = 0 + + async def handler_with_retry(msg, consumer_ref, flow): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise TooManyRequests("rate limited") + # Second call succeeds + + consumer = _make_consumer( + handler=handler_with_retry, + rate_limit_retry_time=0.01, + ) + mock_msg = _make_msg() + consumer.consumer = MagicMock() + + await consumer.handle_one_from_queue(mock_msg) + + assert call_count == 2 + consumer.consumer.acknowledge.assert_called_once_with(mock_msg) + + @pytest.mark.asyncio + async def test_rate_limit_timeout_negative_acks(self): + """If rate limit retries exhaust the timeout, message is negative-acked.""" + async def always_rate_limited(msg, consumer_ref, flow): + raise TooManyRequests("rate limited") + + consumer = _make_consumer( + handler=always_rate_limited, + rate_limit_retry_time=0.01, + rate_limit_timeout=0.05, + ) + mock_msg = _make_msg() + consumer.consumer = MagicMock() + + await consumer.handle_one_from_queue(mock_msg) + + consumer.consumer.negative_acknowledge.assert_called_with(mock_msg) + consumer.consumer.acknowledge.assert_not_called() + + @pytest.mark.asyncio + async def test_non_rate_limit_error_negative_acks_immediately(self): + """Non-TooManyRequests errors should negative-ack immediately (no retry).""" + call_count = 0 + + async def failing_handler(msg, consumer_ref, flow): + nonlocal call_count + call_count += 1 + raise ValueError("bad data") + + consumer = _make_consumer(handler=failing_handler) + mock_msg = _make_msg() + consumer.consumer = MagicMock() + + await consumer.handle_one_from_queue(mock_msg) + + assert call_count == 1 + consumer.consumer.negative_acknowledge.assert_called_once_with(mock_msg) + + @pytest.mark.asyncio + async def test_successful_message_acknowledged(self): + """Successfully processed messages are acknowledged.""" + consumer = _make_consumer(handler=AsyncMock()) + mock_msg = _make_msg() + consumer.consumer = MagicMock() + + await consumer.handle_one_from_queue(mock_msg) + + consumer.consumer.acknowledge.assert_called_once_with(mock_msg) + + +# --------------------------------------------------------------------------- +# Metrics integration +# --------------------------------------------------------------------------- + +class TestMetricsIntegration: + + @pytest.mark.asyncio + async def test_success_metric_on_success(self): + consumer = _make_consumer(handler=AsyncMock()) + mock_msg = _make_msg() + consumer.consumer = MagicMock() + + mock_metrics = MagicMock() + mock_metrics.record_time.return_value.__enter__ = MagicMock() + mock_metrics.record_time.return_value.__exit__ = MagicMock() + consumer.metrics = mock_metrics + + await consumer.handle_one_from_queue(mock_msg) + + mock_metrics.process.assert_called_once_with("success") + + @pytest.mark.asyncio + async def test_error_metric_on_failure(self): + async def failing(msg, c, f): + raise ValueError("fail") + + consumer = _make_consumer(handler=failing) + mock_msg = _make_msg() + consumer.consumer = MagicMock() + + mock_metrics = MagicMock() + consumer.metrics = mock_metrics + + await consumer.handle_one_from_queue(mock_msg) + + mock_metrics.process.assert_called_once_with("error") + + @pytest.mark.asyncio + async def test_rate_limit_metric_on_too_many_requests(self): + call_count = 0 + + async def handler(msg, c, f): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise TooManyRequests("limited") + + consumer = _make_consumer( + handler=handler, + rate_limit_retry_time=0.01, + ) + mock_msg = _make_msg() + consumer.consumer = MagicMock() + + mock_metrics = MagicMock() + mock_metrics.record_time.return_value.__enter__ = MagicMock() + mock_metrics.record_time.return_value.__exit__ = MagicMock(return_value=False) + consumer.metrics = mock_metrics + + await consumer.handle_one_from_queue(mock_msg) + + mock_metrics.rate_limit.assert_called_once() + + +# --------------------------------------------------------------------------- +# Stop / running flag +# --------------------------------------------------------------------------- + +class TestStopBehaviour: + + @pytest.mark.asyncio + async def test_stop_sets_running_false(self): + consumer = _make_consumer() + consumer.running = True + + await consumer.stop() + + assert consumer.running is False + + def test_initial_running_state(self): + consumer = _make_consumer() + assert consumer.running is True diff --git a/tests/unit/test_concurrency/test_dispatcher_semaphore.py b/tests/unit/test_concurrency/test_dispatcher_semaphore.py new file mode 100644 index 00000000..6a1ae8ab --- /dev/null +++ b/tests/unit/test_concurrency/test_dispatcher_semaphore.py @@ -0,0 +1,136 @@ +""" +Tests for MessageDispatcher semaphore-based concurrency enforcement. + +Verifies that the dispatcher limits concurrent message processing to +max_workers via asyncio.Semaphore. +""" + +import asyncio + +import pytest +from unittest.mock import MagicMock, AsyncMock, patch + +from trustgraph.rev_gateway.dispatcher import MessageDispatcher + + +class TestSemaphoreEnforcement: + + @pytest.mark.asyncio + async def test_semaphore_limits_concurrent_processing(self): + """Only max_workers messages should be processed concurrently.""" + max_workers = 2 + dispatcher = MessageDispatcher(max_workers=max_workers) + + concurrent_count = 0 + max_concurrent = 0 + processing_event = asyncio.Event() + + async def slow_process(message): + nonlocal concurrent_count, max_concurrent + concurrent_count += 1 + max_concurrent = max(max_concurrent, concurrent_count) + await asyncio.sleep(0.05) + concurrent_count -= 1 + return {"id": message.get("id"), "response": {"ok": True}} + + dispatcher._process_message = slow_process + + # Launch more tasks than max_workers + messages = [ + {"id": f"msg-{i}", "service": "test", "request": {}} + for i in range(5) + ] + + tasks = [ + asyncio.create_task(dispatcher.handle_message(m)) + for m in messages + ] + + await asyncio.gather(*tasks) + + # At no point should more than max_workers have been active + assert max_concurrent <= max_workers + + @pytest.mark.asyncio + async def test_semaphore_value_matches_max_workers(self): + for n in [1, 5, 20]: + dispatcher = MessageDispatcher(max_workers=n) + assert dispatcher.semaphore._value == n + + @pytest.mark.asyncio + async def test_active_tasks_tracked(self): + """Active tasks should be added/removed during processing.""" + dispatcher = MessageDispatcher(max_workers=5) + + task_was_tracked = False + + original_process = dispatcher._process_message + + async def tracking_process(message): + nonlocal task_was_tracked + # During processing, our task should be in active_tasks + if len(dispatcher.active_tasks) > 0: + task_was_tracked = True + return {"id": message.get("id"), "response": {"ok": True}} + + dispatcher._process_message = tracking_process + + await dispatcher.handle_message( + {"id": "test", "service": "test", "request": {}} + ) + + assert task_was_tracked + # After completion, task should be discarded + assert len(dispatcher.active_tasks) == 0 + + @pytest.mark.asyncio + async def test_semaphore_released_on_error(self): + """Semaphore should be released even if processing raises.""" + dispatcher = MessageDispatcher(max_workers=2) + + async def failing_process(message): + raise RuntimeError("process failed") + + dispatcher._process_message = failing_process + + # Should not deadlock — semaphore must be released on error + with pytest.raises(RuntimeError): + await dispatcher.handle_message( + {"id": "test", "service": "test", "request": {}} + ) + + # Semaphore should be back at max + assert dispatcher.semaphore._value == 2 + + @pytest.mark.asyncio + async def test_single_worker_serializes_processing(self): + """With max_workers=1, messages are processed one at a time.""" + dispatcher = MessageDispatcher(max_workers=1) + + order = [] + + async def ordered_process(message): + msg_id = message["id"] + order.append(f"start-{msg_id}") + await asyncio.sleep(0.02) + order.append(f"end-{msg_id}") + return {"id": msg_id, "response": {"ok": True}} + + dispatcher._process_message = ordered_process + + messages = [{"id": str(i), "service": "t", "request": {}} for i in range(3)] + tasks = [asyncio.create_task(dispatcher.handle_message(m)) for m in messages] + await asyncio.gather(*tasks) + + # With semaphore=1, each message should complete before next starts + # Check that no two "start" entries appear without an intervening "end" + active = 0 + max_active = 0 + for event in order: + if event.startswith("start"): + active += 1 + max_active = max(max_active, active) + elif event.startswith("end"): + active -= 1 + + assert max_active == 1 diff --git a/tests/unit/test_concurrency/test_graph_rag_concurrency.py b/tests/unit/test_concurrency/test_graph_rag_concurrency.py new file mode 100644 index 00000000..8287427b --- /dev/null +++ b/tests/unit/test_concurrency/test_graph_rag_concurrency.py @@ -0,0 +1,268 @@ +""" +Tests for Graph RAG concurrent query execution. + +Covers: execute_batch_triple_queries concurrent task spawning, +exception handling in gather, and result aggregation. +""" + +import asyncio + +import pytest +from unittest.mock import MagicMock, AsyncMock + +from trustgraph.retrieval.graph_rag.graph_rag import Query, LRUCacheWithTTL + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_query( + triples_client=None, + entity_limit=50, + triple_limit=30, + max_subgraph_size=1000, + max_path_length=2, +): + """Create a Query object with mocked rag dependencies.""" + rag = MagicMock() + rag.triples_client = triples_client or AsyncMock() + rag.label_cache = LRUCacheWithTTL() + + query = Query( + rag=rag, + user="test-user", + collection="test-collection", + verbose=False, + entity_limit=entity_limit, + triple_limit=triple_limit, + max_subgraph_size=max_subgraph_size, + max_path_length=max_path_length, + ) + return query + + +def _make_triple(s, p, o): + """Create a simple mock triple.""" + t = MagicMock() + t.s = s + t.p = p + t.o = o + return t + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +class TestBatchTripleQueries: + + @pytest.mark.asyncio + async def test_three_queries_per_entity(self): + """Each entity should generate 3 concurrent queries (s, p, o positions).""" + client = AsyncMock() + client.query_stream = AsyncMock(return_value=[]) + query = _make_query(triples_client=client) + + entities = ["entity-1"] + await query.execute_batch_triple_queries(entities, limit_per_entity=10) + + assert client.query_stream.call_count == 3 + + @pytest.mark.asyncio + async def test_multiple_entities_multiply_queries(self): + """N entities should produce N*3 concurrent queries.""" + client = AsyncMock() + client.query_stream = AsyncMock(return_value=[]) + query = _make_query(triples_client=client) + + entities = ["e1", "e2", "e3"] + await query.execute_batch_triple_queries(entities, limit_per_entity=10) + + assert client.query_stream.call_count == 9 # 3 * 3 + + @pytest.mark.asyncio + async def test_queries_executed_concurrently(self): + """All queries should run concurrently via asyncio.gather.""" + concurrent_count = 0 + max_concurrent = 0 + + async def tracking_query(**kwargs): + nonlocal concurrent_count, max_concurrent + concurrent_count += 1 + max_concurrent = max(max_concurrent, concurrent_count) + await asyncio.sleep(0.02) + concurrent_count -= 1 + return [] + + client = AsyncMock() + client.query_stream = tracking_query + query = _make_query(triples_client=client) + + entities = ["e1", "e2", "e3"] + await query.execute_batch_triple_queries(entities, limit_per_entity=5) + + # All 9 queries should have run concurrently + assert max_concurrent == 9 + + @pytest.mark.asyncio + async def test_results_aggregated(self): + """Results from all queries should be combined into a single list.""" + triple_a = _make_triple("a", "p", "b") + triple_b = _make_triple("c", "p", "d") + + call_count = 0 + + async def alternating_results(**kwargs): + nonlocal call_count + call_count += 1 + if call_count % 2 == 0: + return [triple_a] + return [triple_b] + + client = AsyncMock() + client.query_stream = alternating_results + query = _make_query(triples_client=client) + + result = await query.execute_batch_triple_queries( + ["e1"], limit_per_entity=10 + ) + + # 3 queries, alternating results + assert len(result) == 3 + + @pytest.mark.asyncio + async def test_exception_in_one_query_does_not_block_others(self): + """If one query raises, other results are still collected.""" + good_triple = _make_triple("a", "p", "b") + call_count = 0 + + async def mixed_results(**kwargs): + nonlocal call_count + call_count += 1 + if call_count == 2: + raise RuntimeError("query failed") + return [good_triple] + + client = AsyncMock() + client.query_stream = mixed_results + query = _make_query(triples_client=client) + + result = await query.execute_batch_triple_queries( + ["e1"], limit_per_entity=10 + ) + + # 3 queries: 2 succeed, 1 fails → 2 triples + assert len(result) == 2 + + @pytest.mark.asyncio + async def test_none_results_filtered(self): + """None results from queries should be filtered out.""" + call_count = 0 + + async def sometimes_none(**kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return None + return [_make_triple("a", "p", "b")] + + client = AsyncMock() + client.query_stream = sometimes_none + query = _make_query(triples_client=client) + + result = await query.execute_batch_triple_queries( + ["e1"], limit_per_entity=10 + ) + + # 3 queries: 1 returns None, 2 return triples + assert len(result) == 2 + + @pytest.mark.asyncio + async def test_empty_entities_no_queries(self): + """Empty entity list should produce no queries.""" + client = AsyncMock() + client.query_stream = AsyncMock(return_value=[]) + query = _make_query(triples_client=client) + + result = await query.execute_batch_triple_queries([], limit_per_entity=10) + + assert result == [] + client.query_stream.assert_not_called() + + @pytest.mark.asyncio + async def test_query_params_correct(self): + """Each query should use correct s/p/o positions and params.""" + client = AsyncMock() + client.query_stream = AsyncMock(return_value=[]) + query = _make_query(triples_client=client) + + entities = ["ent-1"] + await query.execute_batch_triple_queries(entities, limit_per_entity=15) + + calls = client.query_stream.call_args_list + assert len(calls) == 3 + + # First call: s=entity, p=None, o=None + assert calls[0].kwargs["s"] == "ent-1" + assert calls[0].kwargs["p"] is None + assert calls[0].kwargs["o"] is None + assert calls[0].kwargs["limit"] == 15 + assert calls[0].kwargs["user"] == "test-user" + assert calls[0].kwargs["collection"] == "test-collection" + assert calls[0].kwargs["batch_size"] == 20 + + # Second call: s=None, p=entity, o=None + assert calls[1].kwargs["s"] is None + assert calls[1].kwargs["p"] == "ent-1" + assert calls[1].kwargs["o"] is None + + # Third call: s=None, p=None, o=entity + assert calls[2].kwargs["s"] is None + assert calls[2].kwargs["p"] is None + assert calls[2].kwargs["o"] == "ent-1" + + +class TestLRUCacheWithTTL: + + def test_put_and_get(self): + cache = LRUCacheWithTTL(max_size=10, ttl=60) + cache.put("key1", "value1") + assert cache.get("key1") == "value1" + + def test_get_missing_returns_none(self): + cache = LRUCacheWithTTL() + assert cache.get("nonexistent") is None + + def test_max_size_eviction(self): + cache = LRUCacheWithTTL(max_size=2, ttl=60) + cache.put("a", 1) + cache.put("b", 2) + cache.put("c", 3) # Should evict "a" + assert cache.get("a") is None + assert cache.get("b") == 2 + assert cache.get("c") == 3 + + def test_lru_order(self): + cache = LRUCacheWithTTL(max_size=2, ttl=60) + cache.put("a", 1) + cache.put("b", 2) + cache.get("a") # Access "a" — now "b" is LRU + cache.put("c", 3) # Should evict "b" + assert cache.get("a") == 1 + assert cache.get("b") is None + assert cache.get("c") == 3 + + def test_ttl_expiration(self): + cache = LRUCacheWithTTL(max_size=10, ttl=0) # TTL=0 means instant expiry + cache.put("key", "value") + # With TTL=0, any time check > 0 means expired + import time + time.sleep(0.01) + assert cache.get("key") is None + + def test_update_existing_key(self): + cache = LRUCacheWithTTL(max_size=10, ttl=60) + cache.put("key", "v1") + cache.put("key", "v2") + assert cache.get("key") == "v2"