test_collectors.pyβ’10.1 kB
"""λ°μ΄ν° μμ§κΈ° ν
μ€νΈ"""
import pytest
import asyncio
from unittest.mock import Mock, patch, AsyncMock
from datetime import datetime, timedelta
from src.collectors.base import BaseCollector
from src.exceptions import DataCollectionError, DataValidationError
class TestBaseCollector:
"""λ² μ΄μ€ λ°μ΄ν° μμ§κΈ° ν
μ€νΈ"""
@pytest.fixture
def collector_config(self):
"""ν
μ€νΈμ© μμ§κΈ° μ€μ """
return {
"name": "test_collector",
"retry_attempts": 3,
"retry_delay": 1,
"timeout": 30,
"batch_size": 100
}
@pytest.fixture
def mock_collector(self, collector_config):
"""ν
μ€νΈμ© Mock μμ§κΈ°"""
class MockCollector(BaseCollector):
async def collect_data(self):
"""Mock λ°μ΄ν° μμ§"""
return [
{"symbol": "005930", "price": 75000, "volume": 1000000},
{"symbol": "000660", "price": 45000, "volume": 500000}
]
def validate_data(self, data):
"""Mock λ°μ΄ν° κ²μ¦"""
for item in data:
if not all(k in item for k in ["symbol", "price", "volume"]):
return False
if item["price"] <= 0 or item["volume"] <= 0:
return False
return True
return MockCollector(collector_config)
def test_base_collector_initialization(self, collector_config):
"""λ² μ΄μ€ μμ§κΈ° μ΄κΈ°ν ν
μ€νΈ"""
class TestCollector(BaseCollector):
async def collect_data(self):
return []
def validate_data(self, data):
return True
collector = TestCollector(collector_config)
assert collector.name == "test_collector"
assert collector.config == collector_config
assert collector.retry_attempts == 3
assert collector.retry_delay == 1
assert collector.timeout == 30
assert collector.batch_size == 100
assert collector.is_running is False
assert collector.last_run is None
assert collector.collected_count == 0
assert collector.error_count == 0
@pytest.mark.asyncio
async def test_successful_data_collection(self, mock_collector):
"""μ±κ³΅μ μΈ λ°μ΄ν° μμ§ ν
μ€νΈ"""
result = await mock_collector.run()
assert result is not None
assert len(result) == 2
assert result[0]["symbol"] == "005930"
assert result[1]["symbol"] == "000660"
assert mock_collector.collected_count == 2
assert mock_collector.error_count == 0
assert mock_collector.last_run is not None
@pytest.mark.asyncio
async def test_data_validation_success(self, mock_collector):
"""λ°μ΄ν° κ²μ¦ μ±κ³΅ ν
μ€νΈ"""
test_data = [
{"symbol": "005930", "price": 75000, "volume": 1000000},
{"symbol": "000660", "price": 45000, "volume": 500000}
]
is_valid = mock_collector.validate_data(test_data)
assert is_valid is True
@pytest.mark.asyncio
async def test_data_validation_failure(self, mock_collector):
"""λ°μ΄ν° κ²μ¦ μ€ν¨ ν
μ€νΈ"""
invalid_data = [
{"symbol": "005930", "price": -1000, "volume": 1000000}, # μμ κ°κ²©
{"symbol": "000660", "price": 45000} # λ³Όλ₯¨ λλ½
]
is_valid = mock_collector.validate_data(invalid_data)
assert is_valid is False
@pytest.mark.asyncio
async def test_retry_on_failure(self, collector_config):
"""μ€ν¨ μ μ¬μλ ν
μ€νΈ"""
call_count = 0
class FailingCollector(BaseCollector):
async def collect_data(self):
nonlocal call_count
call_count += 1
if call_count <= 2: # μ²μ 2λ²μ μ€ν¨
raise Exception("Collection failed")
return [{"symbol": "005930", "price": 75000, "volume": 1000000}]
def validate_data(self, data):
return True
collector = FailingCollector(collector_config)
result = await collector.run()
assert call_count == 3 # 2λ² μ€ν¨ ν 3λ²μ§Έμ μ±κ³΅
assert result is not None
assert len(result) == 1
assert collector.error_count == 2 # 2λ²μ μ€ν¨
assert collector.collected_count == 1
@pytest.mark.asyncio
async def test_max_retry_exceeded(self, collector_config):
"""μ΅λ μ¬μλ μ΄κ³Ό ν
μ€νΈ"""
call_count = 0
class AlwaysFailingCollector(BaseCollector):
async def collect_data(self):
nonlocal call_count
call_count += 1
raise Exception("Always fails")
def validate_data(self, data):
return True
collector = AlwaysFailingCollector(collector_config)
with pytest.raises(DataCollectionError):
await collector.run()
assert call_count == 3 # retry_attemptsλ§νΌ μλ
assert collector.error_count == 3
assert collector.collected_count == 0
@pytest.mark.asyncio
async def test_data_validation_error(self, collector_config):
"""λ°μ΄ν° κ²μ¦ μ€λ₯ ν
μ€νΈ"""
class InvalidDataCollector(BaseCollector):
async def collect_data(self):
return [{"symbol": "005930", "price": -1000, "volume": 1000000}]
def validate_data(self, data):
for item in data:
if item["price"] <= 0:
return False
return True
collector = InvalidDataCollector(collector_config)
with pytest.raises(DataValidationError):
await collector.run()
assert collector.collected_count == 0
assert collector.error_count == 1
@pytest.mark.asyncio
async def test_batch_processing(self, collector_config):
"""λ°°μΉ μ²λ¦¬ ν
μ€νΈ"""
collector_config["batch_size"] = 2 # μμ λ°°μΉ ν¬κΈ°λ‘ ν
μ€νΈ
class BatchCollector(BaseCollector):
async def collect_data(self):
# 5κ° λ°μ΄ν° λ°ν (λ°°μΉ ν¬κΈ° 2λ³΄λ€ νΌ)
return [
{"symbol": f"00{i:04d}", "price": 1000 + i, "volume": 10000 + i}
for i in range(5)
]
def validate_data(self, data):
return True
async def process_batch(self, batch):
"""λ°°μΉ μ²λ¦¬ μ€λ²λΌμ΄λ"""
processed = []
for item in batch:
processed.append({
**item,
"processed_at": datetime.now().isoformat()
})
return processed
collector = BatchCollector(collector_config)
result = await collector.run()
assert len(result) == 5 # μλ³Έ λ°μ΄ν° κ°μ
assert all("processed_at" in item for item in result) # λͺ¨λ μμ΄ν
μ΄ μ²λ¦¬λ¨
assert collector.collected_count == 5
@pytest.mark.asyncio
async def test_concurrent_run_prevention(self, collector_config):
"""λμ μ€ν λ°©μ§ ν
μ€νΈ"""
class SlowCollector(BaseCollector):
async def collect_data(self):
"""λλ¦° Mock λ°μ΄ν° μμ§"""
await asyncio.sleep(0.1) # 0.1μ΄ μ§μ°
return [{"symbol": "005930", "price": 75000, "volume": 1000000}]
def validate_data(self, data):
return True
collector = SlowCollector(collector_config)
# 첫 λ²μ§Έ μ€ν μμ (μλ£λμ§ μμ)
task1 = asyncio.create_task(collector.run())
await asyncio.sleep(0.05) # task1μ΄ μμλλλ‘ μ μ λκΈ°
# λ λ²μ§Έ μ€ν μλ (μ€ν¨ν΄μΌ ν¨)
with pytest.raises(RuntimeError, match="Collector is already running"):
await collector.run()
# 첫 λ²μ§Έ μ€ν μλ£
result1 = await task1
assert result1 is not None
@pytest.mark.asyncio
async def test_stats_collection(self, mock_collector):
"""ν΅κ³ μμ§ ν
μ€νΈ"""
# 첫 λ²μ§Έ μ€ν
await mock_collector.run()
stats1 = mock_collector.get_stats()
assert stats1["total_runs"] == 1
assert stats1["total_collected"] == 2
assert stats1["total_errors"] == 0
assert stats1["success_rate"] == 1.0
# λ λ²μ§Έ μ€ν (μ€ν¨νλλ‘ λ§λ€κΈ°)
original_collect = mock_collector.collect_data
mock_collector.collect_data = AsyncMock(side_effect=Exception("Test error"))
try:
await mock_collector.run()
except DataCollectionError:
pass
# ν΅κ³ νμΈ
stats2 = mock_collector.get_stats()
assert stats2["total_runs"] == 2
assert stats2["total_collected"] == 2 # λ³ν μμ
assert stats2["total_errors"] == 3 # μ¬μλ 3λ²μΌλ‘ μΈν 3κ° μλ¬
assert stats2["success_rate"] == 0.5 # 1/2 = 50%
# μλ ν¨μ 볡μ
mock_collector.collect_data = original_collect
@pytest.mark.asyncio
async def test_cleanup_on_shutdown(self, mock_collector):
"""μ’
λ£ μ μ 리 μμ
ν
μ€νΈ"""
# μ€ν ν μ’
λ£
await mock_collector.run()
await mock_collector.shutdown()
assert mock_collector.is_running is False
# μΆκ°μ μΈ μ 리 μμ
μ΄ μλ€λ©΄ μ¬κΈ°μ νμΈ