BioMCP

"""Tests for request batching functionality.""" import asyncio from unittest.mock import AsyncMock import pytest from biomcp.request_batcher import RequestBatcher @pytest.fixture def mock_batch_func(): """Create a mock batch processing function.""" async def batch_func(params_list): # Simulate processing - return results in same order return [{"result": f"processed_{p['id']}"} for p in params_list] return AsyncMock(side_effect=batch_func) @pytest.mark.asyncio async def test_single_request(mock_batch_func): """Test that a single request is processed correctly.""" batcher = RequestBatcher( batch_func=mock_batch_func, batch_size=10, batch_timeout=0.1 ) result = await batcher.request({"id": 1}) assert result == {"result": "processed_1"} mock_batch_func.assert_called_once_with([{"id": 1}]) @pytest.mark.asyncio async def test_batch_size_trigger(mock_batch_func): """Test that batch is processed when size threshold is reached.""" batcher = RequestBatcher( batch_func=mock_batch_func, batch_size=3, batch_timeout=10 # Long timeout to ensure size triggers first ) # Send requests concurrently tasks = [ batcher.request({"id": 1}), batcher.request({"id": 2}), batcher.request({"id": 3}) ] results = await asyncio.gather(*tasks) assert results == [ {"result": "processed_1"}, {"result": "processed_2"}, {"result": "processed_3"} ] mock_batch_func.assert_called_once_with([ {"id": 1}, {"id": 2}, {"id": 3} ]) @pytest.mark.asyncio async def test_batch_timeout_trigger(mock_batch_func): """Test that batch is processed when timeout expires.""" batcher = RequestBatcher( batch_func=mock_batch_func, batch_size=10, # Large size to ensure timeout triggers first batch_timeout=0.05 ) # Send two requests (less than batch size) task1 = asyncio.create_task(batcher.request({"id": 1})) task2 = asyncio.create_task(batcher.request({"id": 2})) results = await asyncio.gather(task1, task2) assert results == [ {"result": "processed_1"}, {"result": "processed_2"} ] mock_batch_func.assert_called_once() @pytest.mark.asyncio async def test_multiple_batches(mock_batch_func): """Test that multiple batches are processed correctly.""" batcher = RequestBatcher( batch_func=mock_batch_func, batch_size=2, batch_timeout=0.05 ) # Send 5 requests - should create 3 batches tasks = [] for i in range(5): tasks.append(batcher.request({"id": i})) if i < 4: # Small delay between requests await asyncio.sleep(0.01) results = await asyncio.gather(*tasks) assert len(results) == 5 assert all(results[i] == {"result": f"processed_{i}"} for i in range(5)) assert mock_batch_func.call_count >= 2 # At least 2 batches @pytest.mark.asyncio async def test_error_propagation(mock_batch_func): """Test that errors are properly propagated to individual requests.""" # Override mock to raise an error mock_batch_func.side_effect = Exception("Batch processing failed") batcher = RequestBatcher( batch_func=mock_batch_func, batch_size=2, batch_timeout=0.05 ) tasks = [ batcher.request({"id": 1}), batcher.request({"id": 2}) ] # Both requests should receive the error with pytest.raises(Exception, match="Batch processing failed"): await asyncio.gather(*tasks) @pytest.mark.asyncio async def test_partial_error_handling(): """Test handling when batch function returns fewer results.""" async def partial_batch_func(params_list): # Return fewer results than requests return [{"result": "only_one"}] batcher = RequestBatcher( batch_func=partial_batch_func, batch_size=2, batch_timeout=0.05 ) task1 = asyncio.create_task(batcher.request({"id": 1})) task2 = asyncio.create_task(batcher.request({"id": 2})) # First request gets result, second gets error result1 = await task1 assert result1 == {"result": "only_one"} with pytest.raises(Exception, match="No result for request"): await task2 @pytest.mark.asyncio async def test_concurrent_batches(mock_batch_func): """Test that multiple batches can be processed concurrently.""" batcher = RequestBatcher( batch_func=mock_batch_func, batch_size=2, batch_timeout=0.05 ) # Create multiple waves of requests all_tasks = [] for wave in range(3): wave_tasks = [ batcher.request({"id": wave * 10 + i}) for i in range(2) ] all_tasks.extend(wave_tasks) await asyncio.sleep(0.01) # Small delay between waves results = await asyncio.gather(*all_tasks) assert len(results) == 6 assert mock_batch_func.call_count >= 3 @pytest.mark.asyncio async def test_empty_batch_handling(mock_batch_func): """Test that empty batches are handled correctly.""" batcher = RequestBatcher( batch_func=mock_batch_func, batch_size=5, batch_timeout=0.01 ) # Start the timer but don't send any requests timer_task = asyncio.create_task(batcher._batch_timer()) await asyncio.sleep(0.02) timer_task.cancel() # Clean up the task # Should not have called the batch function mock_batch_func.assert_not_called() @pytest.mark.asyncio async def test_request_ordering(mock_batch_func): """Test that results are returned in the correct order.""" # Custom batch function that reverses the order async def reverse_batch_func(params_list): results = [{"result": f"processed_{p['id']}"} for p in params_list] return results # Keep original order mock_batch_func.side_effect = reverse_batch_func batcher = RequestBatcher( batch_func=mock_batch_func, batch_size=4, batch_timeout=0.05 ) tasks = [batcher.request({"id": i}) for i in range(4)] results = await asyncio.gather(*tasks) # Results should match the request order for i in range(4): assert results[i] == {"result": f"processed_{i}"} @pytest.mark.asyncio async def test_batch_func_exception_cleanup(): """Test that batcher state is cleaned up after exceptions.""" call_count = 0 async def failing_batch_func(params_list): nonlocal call_count call_count += 1 if call_count == 1: raise Exception("First batch fails") return [{"result": f"processed_{p['id']}"} for p in params_list] batcher = RequestBatcher( batch_func=failing_batch_func, batch_size=2, batch_timeout=0.05 ) # First batch should fail with pytest.raises(Exception, match="First batch fails"): await asyncio.gather( batcher.request({"id": 1}), batcher.request({"id": 2}) ) # Second batch should succeed results = await asyncio.gather( batcher.request({"id": 3}), batcher.request({"id": 4}) ) assert results == [ {"result": "processed_3"}, {"result": "processed_4"} ]