532 lines
18 KiB
Python
532 lines
18 KiB
Python
|
"""Unit tests for BaseConnector and RateLimiter"""
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
import asyncio
|
||
|
|
import time
|
||
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
||
|
|
|
||
|
|
from sentiment_engine.ingestion.base import RateLimiter, BaseConnector, ConnectorRegistry
|
||
|
|
from sentiment_engine.schemas.config import ConnectorConfig
|
||
|
|
from sentiment_engine.schemas.payload import NormalizedPayload, SourceType, AssetMention
|
||
|
|
|
||
|
|
|
||
|
|
class TestRateLimiter:
|
||
|
|
"""Tests for token bucket rate limiter"""
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def limiter(self):
|
||
|
|
return RateLimiter(rps=10.0, burst=5)
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_initial_burst(self, limiter):
|
||
|
|
"""Should allow burst requests immediately"""
|
||
|
|
for _ in range(5):
|
||
|
|
await limiter.acquire() # Should not block
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_rate_limiting_after_burst(self, limiter):
|
||
|
|
"""Should rate limit after burst is exhausted"""
|
||
|
|
# Exhaust burst
|
||
|
|
for _ in range(5):
|
||
|
|
await limiter.acquire()
|
||
|
|
|
||
|
|
# Next acquire should wait ~0.1s (1/10 rps)
|
||
|
|
start = time.monotonic()
|
||
|
|
await limiter.acquire()
|
||
|
|
elapsed = time.monotonic() - start
|
||
|
|
assert 0.05 < elapsed < 0.3 # ~0.1s with some tolerance
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_token_regeneration(self, limiter):
|
||
|
|
"""Tokens should regenerate over time"""
|
||
|
|
# Exhaust burst
|
||
|
|
for _ in range(5):
|
||
|
|
await limiter.acquire()
|
||
|
|
|
||
|
|
# Wait for tokens to regenerate
|
||
|
|
await asyncio.sleep(0.5) # Should regenerate ~5 tokens at 10 rps
|
||
|
|
|
||
|
|
# Should allow 5 more without waiting
|
||
|
|
start = time.monotonic()
|
||
|
|
for _ in range(5):
|
||
|
|
await limiter.acquire()
|
||
|
|
elapsed = time.monotonic() - start
|
||
|
|
assert elapsed < 0.1 # Should be nearly instant
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_concurrent_access(self, limiter):
|
||
|
|
"""Rate limiter should be thread-safe"""
|
||
|
|
async def acquire_n(n):
|
||
|
|
for _ in range(n):
|
||
|
|
await limiter.acquire()
|
||
|
|
|
||
|
|
await asyncio.gather(acquire_n(3), acquire_n(3), acquire_n(4))
|
||
|
|
# Total 10 acquisitions - should work with burst of 5 + regeneration
|
||
|
|
|
||
|
|
|
||
|
|
class MockConnector(BaseConnector):
|
||
|
|
"""Mock connector for testing"""
|
||
|
|
|
||
|
|
def __init__(self, config: ConnectorConfig, should_fail: bool = False, yield_count: int = 1):
|
||
|
|
super().__init__(config)
|
||
|
|
self.should_fail = should_fail
|
||
|
|
self.yield_count = yield_count
|
||
|
|
self.fetch_called = 0
|
||
|
|
|
||
|
|
async def fetch(self):
|
||
|
|
self.fetch_called += 1
|
||
|
|
if self.should_fail:
|
||
|
|
raise Exception("Simulated fetch error")
|
||
|
|
|
||
|
|
for i in range(self.yield_count):
|
||
|
|
yield NormalizedPayload(
|
||
|
|
source_id=self.config.name,
|
||
|
|
source_type=SourceType.NEWS,
|
||
|
|
source_credibility_base=0.8,
|
||
|
|
ingest_ts=time.time(),
|
||
|
|
publish_ts=time.time(),
|
||
|
|
asset_mentions=[],
|
||
|
|
raw_text=f"Test payload {i}",
|
||
|
|
title=f"Test {i}",
|
||
|
|
url="https://test.com",
|
||
|
|
author="Test",
|
||
|
|
content_length=50,
|
||
|
|
language="en"
|
||
|
|
)
|
||
|
|
|
||
|
|
async def health_check(self) -> bool:
|
||
|
|
return not self.should_fail
|
||
|
|
|
||
|
|
|
||
|
|
class TestBaseConnector:
|
||
|
|
"""Tests for BaseConnector functionality"""
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def config(self):
|
||
|
|
return ConnectorConfig(
|
||
|
|
name="test_connector",
|
||
|
|
source_type="news",
|
||
|
|
poll_interval_seconds=1, # Fast for testing
|
||
|
|
timeout_seconds=5,
|
||
|
|
rate_limit_rps=10.0,
|
||
|
|
rate_limit_rpm=100,
|
||
|
|
rate_limit_burst=5,
|
||
|
|
backoff_base_seconds=0.1,
|
||
|
|
backoff_max_seconds=1.0,
|
||
|
|
backoff_multiplier=2.0,
|
||
|
|
max_concurrent_requests=2,
|
||
|
|
max_latency_ms=1000,
|
||
|
|
min_success_rate=0.5,
|
||
|
|
enabled=True
|
||
|
|
)
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_fetch_yields_payloads(self, config):
|
||
|
|
"""fetch() should yield payloads"""
|
||
|
|
connector = MockConnector(config, should_fail=False, yield_count=3)
|
||
|
|
|
||
|
|
payloads = []
|
||
|
|
async for payload in connector.fetch():
|
||
|
|
payloads.append(payload)
|
||
|
|
|
||
|
|
assert len(payloads) == 3
|
||
|
|
assert connector.fetch_called == 1
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_fetch_raises_on_error(self, config):
|
||
|
|
"""fetch() should raise on error"""
|
||
|
|
connector = MockConnector(config, should_fail=True)
|
||
|
|
|
||
|
|
with pytest.raises(Exception, match="Simulated fetch error"):
|
||
|
|
async for _ in connector.fetch():
|
||
|
|
pass
|
||
|
|
|
||
|
|
assert connector.fetch_called == 1
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_poll_loop_updates_stats_on_success(self, config):
|
||
|
|
"""Poll loop should update stats on successful fetch"""
|
||
|
|
connector = MockConnector(config, should_fail=False, yield_count=3)
|
||
|
|
|
||
|
|
# Manually run one iteration of poll loop logic
|
||
|
|
await connector.rate_limiter.acquire()
|
||
|
|
async with connector.semaphore:
|
||
|
|
async for payload in connector.fetch():
|
||
|
|
connector.stats["total_fetched"] += 1
|
||
|
|
connector.stats["successful"] += 1
|
||
|
|
|
||
|
|
# Success - reset backoff (as done in _run_poll_loop)
|
||
|
|
connector._current_backoff = 0.0
|
||
|
|
connector.stats["consecutive_errors"] = 0
|
||
|
|
connector.stats["last_fetch_ts"] = time.time()
|
||
|
|
|
||
|
|
assert connector.stats["total_fetched"] == 3
|
||
|
|
assert connector.stats["successful"] == 3
|
||
|
|
assert connector.stats["consecutive_errors"] == 0
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_poll_loop_updates_stats_on_error(self, config):
|
||
|
|
"""Poll loop should update stats on fetch error"""
|
||
|
|
connector = MockConnector(config, should_fail=True)
|
||
|
|
|
||
|
|
try:
|
||
|
|
async with connector.semaphore:
|
||
|
|
async for payload in connector.fetch():
|
||
|
|
pass
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
|
||
|
|
# Error handling (as done in _run_poll_loop)
|
||
|
|
connector.stats["errors"] += 1
|
||
|
|
connector.stats["consecutive_errors"] += 1
|
||
|
|
connector._current_backoff = min(
|
||
|
|
connector.backoff_max,
|
||
|
|
max(connector._current_backoff * connector.config.backoff_multiplier, connector.backoff_base)
|
||
|
|
)
|
||
|
|
|
||
|
|
assert connector.stats["errors"] == 1
|
||
|
|
assert connector.stats["consecutive_errors"] == 1
|
||
|
|
assert connector._current_backoff == config.backoff_base_seconds
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_exponential_backoff(self, config):
|
||
|
|
"""Backoff should increase exponentially on consecutive errors"""
|
||
|
|
connector = MockConnector(config, should_fail=True)
|
||
|
|
|
||
|
|
# First error
|
||
|
|
try:
|
||
|
|
async with connector.semaphore:
|
||
|
|
async for _ in connector.fetch(): pass
|
||
|
|
except: pass
|
||
|
|
connector.stats["errors"] += 1
|
||
|
|
connector.stats["consecutive_errors"] += 1
|
||
|
|
connector._current_backoff = min(
|
||
|
|
connector.backoff_max,
|
||
|
|
max(connector._current_backoff * connector.config.backoff_multiplier, connector.backoff_base)
|
||
|
|
)
|
||
|
|
assert connector._current_backoff == config.backoff_base_seconds
|
||
|
|
|
||
|
|
# Second error
|
||
|
|
try:
|
||
|
|
async with connector.semaphore:
|
||
|
|
async for _ in connector.fetch(): pass
|
||
|
|
except: pass
|
||
|
|
connector.stats["errors"] += 1
|
||
|
|
connector.stats["consecutive_errors"] += 1
|
||
|
|
connector._current_backoff = min(
|
||
|
|
connector.backoff_max,
|
||
|
|
max(connector._current_backoff * connector.config.backoff_multiplier, connector.backoff_base)
|
||
|
|
)
|
||
|
|
assert connector._current_backoff == min(config.backoff_max_seconds, config.backoff_base_seconds * 2)
|
||
|
|
|
||
|
|
# Third error
|
||
|
|
try:
|
||
|
|
async with connector.semaphore:
|
||
|
|
async for _ in connector.fetch(): pass
|
||
|
|
except: pass
|
||
|
|
connector.stats["errors"] += 1
|
||
|
|
connector.stats["consecutive_errors"] += 1
|
||
|
|
connector._current_backoff = min(
|
||
|
|
connector.backoff_max,
|
||
|
|
max(connector._current_backoff * connector.config.backoff_multiplier, connector.backoff_base)
|
||
|
|
)
|
||
|
|
assert connector._current_backoff == min(config.backoff_max_seconds, config.backoff_base_seconds * 4)
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_backoff_reset_on_success(self, config):
|
||
|
|
"""Backoff should reset after successful fetch"""
|
||
|
|
connector = MockConnector(config, should_fail=True)
|
||
|
|
|
||
|
|
# Cause an error
|
||
|
|
try:
|
||
|
|
async with connector.semaphore:
|
||
|
|
async for _ in connector.fetch(): pass
|
||
|
|
except: pass
|
||
|
|
connector.stats["errors"] += 1
|
||
|
|
connector.stats["consecutive_errors"] += 1
|
||
|
|
connector._current_backoff = min(
|
||
|
|
connector.backoff_max,
|
||
|
|
max(connector._current_backoff * connector.config.backoff_multiplier, connector.backoff_base)
|
||
|
|
)
|
||
|
|
backoff_after_error = connector._current_backoff
|
||
|
|
|
||
|
|
# Now succeed
|
||
|
|
connector.should_fail = False
|
||
|
|
await connector.rate_limiter.acquire()
|
||
|
|
async with connector.semaphore:
|
||
|
|
async for payload in connector.fetch():
|
||
|
|
connector.stats["total_fetched"] += 1
|
||
|
|
connector.stats["successful"] += 1
|
||
|
|
|
||
|
|
# Success handling
|
||
|
|
connector._current_backoff = 0.0
|
||
|
|
connector.stats["consecutive_errors"] = 0
|
||
|
|
|
||
|
|
assert connector._current_backoff == 0.0
|
||
|
|
assert connector.stats["consecutive_errors"] == 0
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_concurrency_semaphore(self, config):
|
||
|
|
"""Connector should limit concurrent requests"""
|
||
|
|
config.max_concurrent_requests = 1
|
||
|
|
config.poll_interval_seconds = 1
|
||
|
|
|
||
|
|
call_times = []
|
||
|
|
|
||
|
|
class SlowConnector(MockConnector):
|
||
|
|
async def fetch(self):
|
||
|
|
call_times.append(time.monotonic())
|
||
|
|
await asyncio.sleep(0.1) # Simulate slow fetch
|
||
|
|
for payload in super().fetch():
|
||
|
|
yield payload
|
||
|
|
|
||
|
|
connector = SlowConnector(config, should_fail=False, yield_count=1)
|
||
|
|
|
||
|
|
# Start 3 concurrent fetches
|
||
|
|
async def fetch_one():
|
||
|
|
async for p in connector.fetch():
|
||
|
|
return p
|
||
|
|
|
||
|
|
tasks = [asyncio.create_task(fetch_one()) for _ in range(3)]
|
||
|
|
await asyncio.gather(*tasks)
|
||
|
|
|
||
|
|
# With semaphore=1, they should be serialized
|
||
|
|
assert len(call_times) == 3
|
||
|
|
assert call_times[-1] - call_times[0] >= 0.15
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_poll_loop_start_stop(self, config):
|
||
|
|
"""Poll loop should start and stop correctly"""
|
||
|
|
connector = MockConnector(config, should_fail=False, yield_count=1)
|
||
|
|
mock_router = AsyncMock()
|
||
|
|
connector.set_router(mock_router)
|
||
|
|
|
||
|
|
await connector.start()
|
||
|
|
assert connector._running is True
|
||
|
|
assert connector._task is not None
|
||
|
|
|
||
|
|
# Wait for at least one poll cycle
|
||
|
|
await asyncio.sleep(0.2)
|
||
|
|
|
||
|
|
await connector.stop()
|
||
|
|
assert connector._running is False
|
||
|
|
# Task should be cancelled
|
||
|
|
assert connector._task.cancelled()
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_health_check(self, config):
|
||
|
|
"""Health check should reflect connector state"""
|
||
|
|
connector = MockConnector(config, should_fail=False)
|
||
|
|
assert await connector.health_check() is True
|
||
|
|
|
||
|
|
connector.should_fail = True
|
||
|
|
assert await connector.health_check() is False
|
||
|
|
|
||
|
|
|
||
|
|
class TestConnectorRegistry:
|
||
|
|
"""Tests for ConnectorRegistry"""
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def registry(self):
|
||
|
|
return ConnectorRegistry()
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def config(self):
|
||
|
|
return ConnectorConfig(
|
||
|
|
name="test",
|
||
|
|
source_type="news",
|
||
|
|
poll_interval_seconds=1,
|
||
|
|
)
|
||
|
|
|
||
|
|
def test_register_unregister(self, registry, config):
|
||
|
|
connector = MockConnector(config)
|
||
|
|
registry.register(connector)
|
||
|
|
|
||
|
|
assert registry.get("test") is connector
|
||
|
|
assert len(registry.get_all()) == 1
|
||
|
|
assert len(registry.get_enabled()) == 1
|
||
|
|
|
||
|
|
registry.unregister("test")
|
||
|
|
assert registry.get("test") is None
|
||
|
|
assert len(registry.get_all()) == 0
|
||
|
|
|
||
|
|
def test_disabled_connector_not_in_enabled(self, registry, config):
|
||
|
|
config.enabled = False
|
||
|
|
connector = MockConnector(config)
|
||
|
|
registry.register(connector)
|
||
|
|
|
||
|
|
assert len(registry.get_all()) == 1
|
||
|
|
assert len(registry.get_enabled()) == 0
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_start_stop_all(self, registry, config):
|
||
|
|
connector1 = MockConnector(config)
|
||
|
|
connector2 = MockConnector(config)
|
||
|
|
|
||
|
|
registry.register(connector1)
|
||
|
|
registry.register(connector2)
|
||
|
|
|
||
|
|
mock_router = AsyncMock()
|
||
|
|
registry.set_router(mock_router)
|
||
|
|
|
||
|
|
await registry.start_all()
|
||
|
|
assert connector1._running
|
||
|
|
assert connector2._running
|
||
|
|
|
||
|
|
await registry.stop_all()
|
||
|
|
assert not connector1._running
|
||
|
|
assert not connector2._running
|
||
|
|
|
||
|
|
def test_set_router(self, registry):
|
||
|
|
mock_router = MagicMock()
|
||
|
|
registry.set_router(mock_router)
|
||
|
|
assert registry._router is mock_router
|
||
|
|
|
||
|
|
|
||
|
|
# Pairwise tests - multiple connectors interacting
|
||
|
|
class TestPairwiseConnectors:
|
||
|
|
"""Tests for multiple connectors running together"""
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_two_connectors_independent(self):
|
||
|
|
"""Two connectors should operate independently"""
|
||
|
|
config1 = ConnectorConfig(
|
||
|
|
name="connector1", source_type="news", poll_interval_seconds=1,
|
||
|
|
rate_limit_rps=10, rate_limit_burst=5
|
||
|
|
)
|
||
|
|
config2 = ConnectorConfig(
|
||
|
|
name="connector2", source_type="news", poll_interval_seconds=1,
|
||
|
|
rate_limit_rps=10, rate_limit_burst=5
|
||
|
|
)
|
||
|
|
|
||
|
|
conn1 = MockConnector(config1, yield_count=3)
|
||
|
|
conn2 = MockConnector(config2, yield_count=2)
|
||
|
|
|
||
|
|
# Run concurrently
|
||
|
|
results1 = []
|
||
|
|
results2 = []
|
||
|
|
|
||
|
|
async def collect1():
|
||
|
|
async for p in conn1.fetch():
|
||
|
|
results1.append(p)
|
||
|
|
|
||
|
|
async def collect2():
|
||
|
|
async for p in conn2.fetch():
|
||
|
|
results2.append(p)
|
||
|
|
|
||
|
|
await asyncio.gather(collect1(), collect2())
|
||
|
|
|
||
|
|
assert len(results1) == 3
|
||
|
|
assert len(results2) == 2
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_registry_routes_to_router(self):
|
||
|
|
"""Registry should route payloads to router"""
|
||
|
|
registry = ConnectorRegistry()
|
||
|
|
mock_router = AsyncMock()
|
||
|
|
registry.set_router(mock_router)
|
||
|
|
|
||
|
|
config = ConnectorConfig(name="test", source_type="news", poll_interval_seconds=1)
|
||
|
|
connector = MockConnector(config, yield_count=1)
|
||
|
|
registry.register(connector)
|
||
|
|
|
||
|
|
payload = NormalizedPayload(
|
||
|
|
source_id="test", source_type=SourceType.NEWS,
|
||
|
|
source_credibility_base=0.8, ingest_ts=time.time(),
|
||
|
|
publish_ts=time.time(), asset_mentions=[],
|
||
|
|
raw_text="test", title="test", url="https://test.com",
|
||
|
|
author="test", content_length=10, language="en"
|
||
|
|
)
|
||
|
|
|
||
|
|
await registry.route_payload(payload)
|
||
|
|
mock_router.route.assert_called_once_with(payload)
|
||
|
|
|
||
|
|
|
||
|
|
# E2E-style test for full connector lifecycle
|
||
|
|
class TestConnectorLifecycle:
|
||
|
|
"""Full lifecycle tests for connectors"""
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_full_lifecycle(self):
|
||
|
|
"""Test complete connector lifecycle: start -> fetch -> stats -> stop"""
|
||
|
|
config = ConnectorConfig(
|
||
|
|
name="lifecycle_test",
|
||
|
|
source_type="news",
|
||
|
|
poll_interval_seconds=1, # Must be int >= 1
|
||
|
|
rate_limit_rps=100,
|
||
|
|
rate_limit_burst=10,
|
||
|
|
timeout_seconds=1,
|
||
|
|
backoff_base_seconds=0.1,
|
||
|
|
backoff_max_seconds=1.0,
|
||
|
|
backoff_multiplier=2.0,
|
||
|
|
max_concurrent_requests=2,
|
||
|
|
max_latency_ms=1000,
|
||
|
|
min_success_rate=0.5,
|
||
|
|
enabled=True
|
||
|
|
)
|
||
|
|
|
||
|
|
connector = MockConnector(config, yield_count=2)
|
||
|
|
mock_router = AsyncMock()
|
||
|
|
connector.set_router(mock_router)
|
||
|
|
|
||
|
|
# Start
|
||
|
|
await connector.start()
|
||
|
|
assert connector._running
|
||
|
|
|
||
|
|
# Let it run a few cycles
|
||
|
|
await asyncio.sleep(0.3)
|
||
|
|
|
||
|
|
# Check stats
|
||
|
|
stats = connector.get_stats()
|
||
|
|
assert stats["total_fetched"] > 0
|
||
|
|
assert stats["success_rate"] == 1.0
|
||
|
|
|
||
|
|
# Stop
|
||
|
|
await connector.stop()
|
||
|
|
assert not connector._running
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_lifecycle_with_errors(self):
|
||
|
|
"""Lifecycle with mixed success/failure"""
|
||
|
|
config = ConnectorConfig(
|
||
|
|
name="error_test",
|
||
|
|
source_type="news",
|
||
|
|
poll_interval_seconds=1,
|
||
|
|
rate_limit_rps=100,
|
||
|
|
backoff_base_seconds=0.1,
|
||
|
|
backoff_max_seconds=1.0,
|
||
|
|
enabled=True
|
||
|
|
)
|
||
|
|
|
||
|
|
connector = MockConnector(config, should_fail=False, yield_count=1)
|
||
|
|
mock_router = AsyncMock()
|
||
|
|
connector.set_router(mock_router)
|
||
|
|
|
||
|
|
await connector.start()
|
||
|
|
|
||
|
|
# Let it succeed a few times
|
||
|
|
await asyncio.sleep(0.2)
|
||
|
|
|
||
|
|
# Cause failures
|
||
|
|
connector.should_fail = True
|
||
|
|
await asyncio.sleep(0.3)
|
||
|
|
|
||
|
|
# Should have backoff
|
||
|
|
assert connector._current_backoff > 0
|
||
|
|
assert connector.stats["errors"] > 0
|
||
|
|
assert connector.stats["consecutive_errors"] > 0
|
||
|
|
|
||
|
|
# Recover
|
||
|
|
connector.should_fail = False
|
||
|
|
await asyncio.sleep(0.2)
|
||
|
|
|
||
|
|
# Backoff should reset
|
||
|
|
assert connector._current_backoff == 0
|
||
|
|
assert connector.stats["consecutive_errors"] == 0
|
||
|
|
|
||
|
|
await connector.stop()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
pytest.main([__file__, "-v"])
|