224 lines
7.8 KiB
Python
224 lines
7.8 KiB
Python
"""
|
|
Tests for RabbitMQ client.
|
|
"""
|
|
|
|
import pytest
|
|
import json
|
|
from unittest.mock import Mock, patch, MagicMock
|
|
import pika
|
|
|
|
from app.queues.rabbitmq_client import RabbitMQClient, create_rabbitmq_client
|
|
|
|
|
|
class TestRabbitMQClient:
|
|
"""Tests for RabbitMQClient class."""
|
|
|
|
def test_initialization(self):
|
|
"""Test RabbitMQ client initialization."""
|
|
client = RabbitMQClient(
|
|
rabbitmq_url="amqp://guest:guest@localhost:5672",
|
|
prefetch_count=1
|
|
)
|
|
|
|
assert client.rabbitmq_url == "amqp://guest:guest@localhost:5672"
|
|
assert client.prefetch_count == 1
|
|
assert client.connection is None
|
|
assert client.channel is None
|
|
assert client.queues == {
|
|
'scraping_tasks': 'scraping_tasks',
|
|
'sentiment_analysis_tasks': 'sentiment_analysis_tasks',
|
|
'energy_calculation_tasks': 'energy_calculation_tasks',
|
|
'results': 'results'
|
|
}
|
|
|
|
@patch('pika.BlockingConnection')
|
|
def test_connect_success(self, mock_connection):
|
|
"""Test successful connection to RabbitMQ."""
|
|
# Setup mocks
|
|
mock_conn_instance = Mock()
|
|
mock_channel_instance = Mock()
|
|
mock_connection.return_value = mock_conn_instance
|
|
mock_conn_instance.channel.return_value = mock_channel_instance
|
|
|
|
# Create client and connect
|
|
client = RabbitMQClient()
|
|
client.connect()
|
|
|
|
# Verify connection and channel created
|
|
mock_connection.assert_called_once()
|
|
mock_conn_instance.channel.assert_called_once()
|
|
mock_channel_instance.basic_qos.assert_called_once_with(prefetch_count=1)
|
|
|
|
# Verify queues declared
|
|
assert mock_channel_instance.queue_declare.call_count == 4
|
|
|
|
# Verify client state
|
|
assert client.connection == mock_conn_instance
|
|
assert client.channel == mock_channel_instance
|
|
|
|
@patch('pika.BlockingConnection')
|
|
def test_publish_message(self, mock_connection):
|
|
"""Test publishing a message to a queue."""
|
|
# Setup mocks
|
|
mock_conn_instance = Mock()
|
|
mock_channel_instance = Mock()
|
|
mock_connection.return_value = mock_conn_instance
|
|
mock_conn_instance.channel.return_value = mock_channel_instance
|
|
|
|
# Create client and connect
|
|
client = RabbitMQClient()
|
|
client.connect()
|
|
|
|
# Publish message
|
|
client.publish_message(
|
|
queue_name='scraping_tasks',
|
|
data={'match_id': 123, 'source': 'twitter'},
|
|
event_type='scraping.task.created'
|
|
)
|
|
|
|
# Verify publish called
|
|
mock_channel_instance.basic_publish.assert_called_once()
|
|
call_args = mock_channel_instance.basic_publish.call_args
|
|
|
|
# Check routing key
|
|
assert call_args[1]['routing_key'] == 'scraping_tasks'
|
|
|
|
# Check message body is valid JSON
|
|
message_body = call_args[1]['body']
|
|
message = json.loads(message_body)
|
|
|
|
assert message['event'] == 'scraping.task.created'
|
|
assert message['version'] == '1.0'
|
|
assert message['data']['match_id'] == 123
|
|
assert message['data']['source'] == 'twitter'
|
|
assert 'timestamp' in message
|
|
|
|
# Check message properties
|
|
properties = call_args[1]['properties']
|
|
assert properties.delivery_mode == 2 # Persistent
|
|
|
|
@patch('pika.BlockingConnection')
|
|
def test_consume_messages(self, mock_connection):
|
|
"""Test consuming messages from a queue."""
|
|
# Setup mocks
|
|
mock_conn_instance = Mock()
|
|
mock_channel_instance = Mock()
|
|
mock_connection.return_value = mock_conn_instance
|
|
mock_conn_instance.channel.return_value = mock_channel_instance
|
|
|
|
# Create client and connect
|
|
client = RabbitMQClient()
|
|
client.connect()
|
|
|
|
# Mock start_consuming to avoid blocking
|
|
mock_channel_instance.start_consuming = Mock(side_effect=KeyboardInterrupt())
|
|
|
|
# Define callback
|
|
callback = Mock()
|
|
|
|
# Consume messages
|
|
try:
|
|
client.consume_messages(queue_name='scraping_tasks', callback=callback)
|
|
except KeyboardInterrupt:
|
|
pass
|
|
|
|
# Verify consume called
|
|
mock_channel_instance.basic_consume.assert_called_once()
|
|
call_args = mock_channel_instance.basic_consume.call_args
|
|
|
|
assert call_args[1]['queue'] == 'scraping_tasks'
|
|
assert call_args[1]['auto_ack'] is False
|
|
assert call_args[1]['on_message_callback'] == callback
|
|
|
|
# Verify start_consuming called
|
|
mock_channel_instance.start_consuming.assert_called_once()
|
|
|
|
@patch('pika.BlockingConnection')
|
|
def test_ack_message(self, mock_connection):
|
|
"""Test acknowledging a message."""
|
|
# Setup mocks
|
|
mock_conn_instance = Mock()
|
|
mock_channel_instance = Mock()
|
|
mock_connection.return_value = mock_conn_instance
|
|
mock_conn_instance.channel.return_value = mock_channel_instance
|
|
|
|
# Create client and connect
|
|
client = RabbitMQClient()
|
|
client.connect()
|
|
|
|
# Acknowledge message
|
|
client.ack_message(delivery_tag=123)
|
|
|
|
# Verify ack called
|
|
mock_channel_instance.basic_ack.assert_called_once_with(delivery_tag=123)
|
|
|
|
@patch('pika.BlockingConnection')
|
|
def test_reject_message(self, mock_connection):
|
|
"""Test rejecting a message."""
|
|
# Setup mocks
|
|
mock_conn_instance = Mock()
|
|
mock_channel_instance = Mock()
|
|
mock_connection.return_value = mock_conn_instance
|
|
mock_conn_instance.channel.return_value = mock_channel_instance
|
|
|
|
# Create client and connect
|
|
client = RabbitMQClient()
|
|
client.connect()
|
|
|
|
# Reject message without requeue
|
|
client.reject_message(delivery_tag=123, requeue=False)
|
|
|
|
# Verify reject called
|
|
mock_channel_instance.basic_reject.assert_called_once_with(
|
|
delivery_tag=123,
|
|
requeue=False
|
|
)
|
|
|
|
@patch('pika.BlockingConnection')
|
|
def test_close_connection(self, mock_connection):
|
|
"""Test closing connection to RabbitMQ."""
|
|
# Setup mocks
|
|
mock_conn_instance = Mock()
|
|
mock_channel_instance = Mock()
|
|
mock_connection.return_value = mock_conn_instance
|
|
mock_conn_instance.channel.return_value = mock_channel_instance
|
|
|
|
# Create client and connect
|
|
client = RabbitMQClient()
|
|
client.connect()
|
|
|
|
# Close connection
|
|
client.close()
|
|
|
|
# Verify channel and connection closed
|
|
mock_channel_instance.close.assert_called_once()
|
|
mock_conn_instance.close.assert_called_once()
|
|
|
|
|
|
class TestCreateRabbitMQClient:
|
|
"""Tests for create_rabbitmq_client factory function."""
|
|
|
|
def test_create_with_defaults(self):
|
|
"""Test creating client with default parameters."""
|
|
client = create_rabbitmq_client()
|
|
|
|
assert client.rabbitmq_url == "amqp://guest:guest@localhost:5672"
|
|
assert client.prefetch_count == 1
|
|
assert isinstance(client, RabbitMQClient)
|
|
|
|
def test_create_with_custom_url(self):
|
|
"""Test creating client with custom URL."""
|
|
client = create_rabbitmq_client(
|
|
rabbitmq_url="amqp://user:pass@remote:5672"
|
|
)
|
|
|
|
assert client.rabbitmq_url == "amqp://user:pass@remote:5672"
|
|
assert client.prefetch_count == 1
|
|
|
|
def test_create_with_custom_prefetch(self):
|
|
"""Test creating client with custom prefetch count."""
|
|
client = create_rabbitmq_client(prefetch_count=5)
|
|
|
|
assert client.rabbitmq_url == "amqp://guest:guest@localhost:5672"
|
|
assert client.prefetch_count == 5
|