282 lines
9.3 KiB
Python
282 lines
9.3 KiB
Python
"""
|
|
Tests for energy calculation worker.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock, patch
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.workers.energy_worker import (
|
|
EnergyWorker,
|
|
create_energy_worker
|
|
)
|
|
|
|
|
|
class TestEnergyWorker:
|
|
"""Tests for EnergyWorker class."""
|
|
|
|
def test_initialization(self):
|
|
"""Test energy worker initialization."""
|
|
worker = EnergyWorker()
|
|
|
|
# No specific initialization required for energy worker
|
|
assert worker is not None
|
|
|
|
@patch('app.workers.energy_worker.calculate_and_store_energy_score')
|
|
@patch('app.workers.energy_worker.get_energy_score_by_match_and_team')
|
|
def test_execute_energy_calculation_task_new(
|
|
self,
|
|
mock_get_score,
|
|
mock_calculate_store
|
|
):
|
|
"""Test executing energy calculation for new score."""
|
|
# Create worker
|
|
worker = EnergyWorker()
|
|
|
|
# Mock database session
|
|
mock_db = Mock(spec=Session)
|
|
|
|
# Mock no existing score
|
|
mock_get_score.return_value = None
|
|
|
|
# Mock calculation
|
|
mock_energy_score = Mock()
|
|
mock_energy_score.score = 75.5
|
|
mock_energy_score.confidence = 0.82
|
|
mock_energy_score.sources_used = ['twitter', 'reddit']
|
|
mock_energy_score.twitter_score = 0.6
|
|
mock_energy_score.reddit_score = 0.7
|
|
mock_energy_score.rss_score = 0.5
|
|
mock_energy_score.temporal_factor = 1.0
|
|
mock_calculate_store.return_value = mock_energy_score
|
|
|
|
# Execute task
|
|
task = {
|
|
'match_id': 123,
|
|
'team_id': 456,
|
|
'twitter_sentiments': [
|
|
{'compound': 0.6, 'sentiment': 'positive'},
|
|
{'compound': 0.7, 'sentiment': 'positive'}
|
|
],
|
|
'reddit_sentiments': [
|
|
{'compound': 0.7, 'sentiment': 'positive'},
|
|
{'compound': 0.8, 'sentiment': 'positive'}
|
|
],
|
|
'rss_sentiments': [
|
|
{'compound': 0.5, 'sentiment': 'positive'}
|
|
],
|
|
'tweets_with_timestamps': []
|
|
}
|
|
|
|
result = worker.execute_energy_calculation_task(task, mock_db)
|
|
|
|
# Verify calculation called
|
|
mock_calculate_store.assert_called_once()
|
|
|
|
# Verify result
|
|
assert result['energy_score'] == 75.5
|
|
assert result['confidence'] == 0.82
|
|
assert result['sources_used'] == ['twitter', 'reddit']
|
|
assert result['status'] == 'success'
|
|
assert result['metadata']['match_id'] == 123
|
|
assert result['metadata']['team_id'] == 456
|
|
assert result['metadata']['twitter_score'] == 0.6
|
|
assert result['metadata']['reddit_score'] == 0.7
|
|
assert result['metadata']['rss_score'] == 0.5
|
|
assert result['metadata']['temporal_factor'] == 1.0
|
|
|
|
@patch('app.workers.energy_worker.get_energy_score_by_match_and_team')
|
|
def test_execute_energy_calculation_task_existing(
|
|
self,
|
|
mock_get_score
|
|
):
|
|
"""Test executing energy calculation for existing score."""
|
|
# Create worker
|
|
worker = EnergyWorker()
|
|
|
|
# Mock database session
|
|
mock_db = Mock(spec=Session)
|
|
|
|
# Mock existing score
|
|
mock_existing_score = Mock()
|
|
mock_existing_score.score = 72.0
|
|
mock_existing_score.confidence = 0.78
|
|
mock_existing_score.sources_used = ['twitter', 'reddit']
|
|
mock_get_score.return_value = mock_existing_score
|
|
|
|
# Execute task
|
|
task = {
|
|
'match_id': 123,
|
|
'team_id': 456,
|
|
'twitter_sentiments': [],
|
|
'reddit_sentiments': [],
|
|
'rss_sentiments': [],
|
|
'tweets_with_timestamps': []
|
|
}
|
|
|
|
result = worker.execute_energy_calculation_task(task, mock_db)
|
|
|
|
# Verify result from existing score
|
|
assert result['energy_score'] == 72.0
|
|
assert result['confidence'] == 0.78
|
|
assert result['sources_used'] == ['twitter', 'reddit']
|
|
assert result['status'] == 'success'
|
|
assert result['metadata']['updated_existing'] is True
|
|
|
|
@patch('app.workers.energy_worker.calculate_and_store_energy_score')
|
|
@patch('app.workers.energy_worker.get_energy_score_by_match_and_team')
|
|
def test_execute_energy_calculation_task_error_handling(
|
|
self,
|
|
mock_get_score,
|
|
mock_calculate_store
|
|
):
|
|
"""Test error handling in energy calculation."""
|
|
worker = EnergyWorker()
|
|
mock_db = Mock(spec=Session)
|
|
|
|
# Mock no existing score
|
|
mock_get_score.return_value = None
|
|
|
|
# Mock calculation error
|
|
mock_calculate_store.side_effect = Exception("Calculation error")
|
|
|
|
# Execute task
|
|
task = {
|
|
'match_id': 123,
|
|
'team_id': 456,
|
|
'twitter_sentiments': [],
|
|
'reddit_sentiments': [],
|
|
'rss_sentiments': [],
|
|
'tweets_with_timestamps': []
|
|
}
|
|
|
|
result = worker.execute_energy_calculation_task(task, mock_db)
|
|
|
|
# Verify error handling
|
|
assert result['energy_score'] == 0.0
|
|
assert result['confidence'] == 0.0
|
|
assert result['sources_used'] == []
|
|
assert result['status'] == 'error'
|
|
assert 'error' in result
|
|
|
|
@patch('app.workers.energy_worker.calculate_energy_score')
|
|
def test_calculate_mock_energy(self, mock_calculate):
|
|
"""Test mock energy calculation (for testing)."""
|
|
worker = EnergyWorker()
|
|
|
|
# Mock calculation
|
|
mock_calculate.return_value = {
|
|
'score': 75.5,
|
|
'confidence': 0.82,
|
|
'sources_used': ['twitter', 'reddit']
|
|
}
|
|
|
|
# Calculate mock energy
|
|
twitter_sentiments = [
|
|
{'compound': 0.6, 'sentiment': 'positive'}
|
|
]
|
|
reddit_sentiments = [
|
|
{'compound': 0.7, 'sentiment': 'positive'}
|
|
]
|
|
|
|
result = worker.calculate_mock_energy(
|
|
twitter_sentiments=twitter_sentiments,
|
|
reddit_sentiments=reddit_sentiments
|
|
)
|
|
|
|
# Verify calculation called
|
|
mock_calculate.assert_called_once_with(
|
|
match_id=0,
|
|
team_id=0,
|
|
twitter_sentiments=twitter_sentiments,
|
|
reddit_sentiments=reddit_sentiments,
|
|
rss_sentiments=[],
|
|
tweets_with_timestamps=[]
|
|
)
|
|
|
|
# Verify result
|
|
assert result['energy_score'] == 75.5
|
|
assert result['confidence'] == 0.82
|
|
assert result['sources_used'] == ['twitter', 'reddit']
|
|
assert result['status'] == 'success'
|
|
|
|
@patch('app.workers.energy_worker.calculate_energy_score')
|
|
def test_calculate_mock_energy_error(self, mock_calculate):
|
|
"""Test error handling in mock energy calculation."""
|
|
worker = EnergyWorker()
|
|
|
|
# Mock calculation error
|
|
mock_calculate.side_effect = Exception("Calculation error")
|
|
|
|
# Calculate mock energy
|
|
result = worker.calculate_mock_energy(
|
|
twitter_sentiments=[],
|
|
reddit_sentiments=[]
|
|
)
|
|
|
|
# Verify error handling
|
|
assert result['energy_score'] == 0.0
|
|
assert result['confidence'] == 0.0
|
|
assert result['sources_used'] == []
|
|
assert result['status'] == 'error'
|
|
assert 'error' in result
|
|
|
|
@patch('app.workers.energy_worker.calculate_energy_score')
|
|
def test_calculate_mock_energy_with_all_sources(self, mock_calculate):
|
|
"""Test mock energy calculation with all sources."""
|
|
worker = EnergyWorker()
|
|
|
|
# Mock calculation
|
|
mock_calculate.return_value = {
|
|
'score': 80.0,
|
|
'confidence': 0.85,
|
|
'sources_used': ['twitter', 'reddit', 'rss']
|
|
}
|
|
|
|
# Calculate mock energy with all sources
|
|
twitter_sentiments = [
|
|
{'compound': 0.6, 'sentiment': 'positive'}
|
|
]
|
|
reddit_sentiments = [
|
|
{'compound': 0.7, 'sentiment': 'positive'}
|
|
]
|
|
rss_sentiments = [
|
|
{'compound': 0.5, 'sentiment': 'positive'}
|
|
]
|
|
tweets_with_timestamps = [
|
|
{'timestamp': '2026-01-17T10:00:00', 'text': 'Great match!'}
|
|
]
|
|
|
|
result = worker.calculate_mock_energy(
|
|
twitter_sentiments=twitter_sentiments,
|
|
reddit_sentiments=reddit_sentiments,
|
|
rss_sentiments=rss_sentiments,
|
|
tweets_with_timestamps=tweets_with_timestamps
|
|
)
|
|
|
|
# Verify calculation called with all sources
|
|
mock_calculate.assert_called_once_with(
|
|
match_id=0,
|
|
team_id=0,
|
|
twitter_sentiments=twitter_sentiments,
|
|
reddit_sentiments=reddit_sentiments,
|
|
rss_sentiments=rss_sentiments,
|
|
tweets_with_timestamps=tweets_with_timestamps
|
|
)
|
|
|
|
# Verify result
|
|
assert result['energy_score'] == 80.0
|
|
assert result['confidence'] == 0.85
|
|
assert result['sources_used'] == ['twitter', 'reddit', 'rss']
|
|
assert result['status'] == 'success'
|
|
|
|
|
|
class TestCreateEnergyWorker:
|
|
"""Tests for create_energy_worker factory function."""
|
|
|
|
def test_create_energy_worker(self):
|
|
"""Test creating an energy worker."""
|
|
worker = create_energy_worker()
|
|
|
|
assert isinstance(worker, EnergyWorker)
|