147 lines
4.9 KiB
Python
147 lines
4.9 KiB
Python
"""
|
|
Tests for RabbitMQ message consumers.
|
|
"""
|
|
|
|
import pytest
|
|
import json
|
|
from unittest.mock import Mock, patch, MagicMock
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.queues.consumers import (
|
|
consume_scraping_tasks,
|
|
consume_sentiment_analysis_tasks,
|
|
consume_energy_calculation_tasks,
|
|
consume_results
|
|
)
|
|
|
|
|
|
class TestConsumeScrapingTasks:
|
|
"""Tests for consume_scraping_tasks function."""
|
|
|
|
@patch('app.queues.consumers.client')
|
|
@patch('app.queues.consumers.consume_messages')
|
|
def test_consume_scraping_tasks_success(self, mock_consume, mock_client):
|
|
"""Test successful scraping task consumption."""
|
|
# Setup mocks
|
|
mock_db_session = Mock(spec=Session)
|
|
mock_db_factory = Mock(return_value=mock_db_session)
|
|
mock_callback = Mock(return_value={'collected_count': 50})
|
|
|
|
# Call consume
|
|
consume_scraping_tasks(
|
|
client=mock_client,
|
|
callback=mock_callback,
|
|
db_session_factory=mock_db_factory
|
|
)
|
|
|
|
# Verify consume_messages called
|
|
mock_consume.assert_called_once()
|
|
call_args = mock_consume.call_args
|
|
|
|
assert call_args[0][0] == 'scraping_tasks'
|
|
assert call_args[1]['callback'] is not None
|
|
|
|
@patch('app.queues.consumers.client')
|
|
@patch('app.queues.consumers.consume_messages')
|
|
def test_consume_scraping_tasks_error_handling(self, mock_consume, mock_client):
|
|
"""Test error handling in scraping task consumption."""
|
|
mock_db_session = Mock(spec=Session)
|
|
mock_db_factory = Mock(return_value=mock_db_session)
|
|
mock_callback = Mock(side_effect=Exception("Processing error"))
|
|
|
|
# Get the callback function
|
|
def capture_callback(ch, method, properties, body):
|
|
message = json.loads(body)
|
|
mock_callback(message.get('data', {}), mock_db_session)
|
|
|
|
# Call consume
|
|
consume_scraping_tasks(
|
|
client=mock_client,
|
|
callback=mock_callback,
|
|
db_session_factory=mock_db_factory
|
|
)
|
|
|
|
# Verify consume_messages called
|
|
mock_consume.assert_called_once()
|
|
|
|
|
|
class TestConsumeSentimentAnalysisTasks:
|
|
"""Tests for consume_sentiment_analysis_tasks function."""
|
|
|
|
@patch('app.queues.consumers.client')
|
|
@patch('app.queues.consumers.consume_messages')
|
|
def test_consume_sentiment_analysis_tasks_success(self, mock_consume, mock_client):
|
|
"""Test successful sentiment analysis task consumption."""
|
|
# Setup mocks
|
|
mock_db_session = Mock(spec=Session)
|
|
mock_db_factory = Mock(return_value=mock_db_session)
|
|
mock_callback = Mock(return_value={'analyzed_count': 100})
|
|
|
|
# Call consume
|
|
consume_sentiment_analysis_tasks(
|
|
client=mock_client,
|
|
callback=mock_callback,
|
|
db_session_factory=mock_db_factory
|
|
)
|
|
|
|
# Verify consume_messages called
|
|
mock_consume.assert_called_once()
|
|
call_args = mock_consume.call_args
|
|
|
|
assert call_args[0][0] == 'sentiment_analysis_tasks'
|
|
assert call_args[1]['callback'] is not None
|
|
|
|
|
|
class TestConsumeEnergyCalculationTasks:
|
|
"""Tests for consume_energy_calculation_tasks function."""
|
|
|
|
@patch('app.queues.consumers.client')
|
|
@patch('app.queues.consumers.consume_messages')
|
|
def test_consume_energy_calculation_tasks_success(self, mock_consume, mock_client):
|
|
"""Test successful energy calculation task consumption."""
|
|
# Setup mocks
|
|
mock_db_session = Mock(spec=Session)
|
|
mock_db_factory = Mock(return_value=mock_db_session)
|
|
mock_callback = Mock(return_value={
|
|
'energy_score': 75.5,
|
|
'confidence': 0.82
|
|
})
|
|
|
|
# Call consume
|
|
consume_energy_calculation_tasks(
|
|
client=mock_client,
|
|
callback=mock_callback,
|
|
db_session_factory=mock_db_factory
|
|
)
|
|
|
|
# Verify consume_messages called
|
|
mock_consume.assert_called_once()
|
|
call_args = mock_consume.call_args
|
|
|
|
assert call_args[0][0] == 'energy_calculation_tasks'
|
|
assert call_args[1]['callback'] is not None
|
|
|
|
|
|
class TestConsumeResults:
|
|
"""Tests for consume_results function."""
|
|
|
|
@patch('app.queues.consumers.client')
|
|
@patch('app.queues.consumers.consume_messages')
|
|
def test_consume_results_success(self, mock_consume, mock_client):
|
|
"""Test successful result consumption."""
|
|
# Setup mocks
|
|
mock_callback = Mock()
|
|
|
|
# Call consume
|
|
consume_results(
|
|
client=mock_client,
|
|
callback=mock_callback
|
|
)
|
|
|
|
# Verify consume_messages called
|
|
mock_consume.assert_called_once()
|
|
call_args = mock_consume.call_args
|
|
|
|
assert call_args[0][0] == 'results'
|
|
assert call_args[1]['callback'] is not None
|