312 lines
12 KiB
Python
312 lines
12 KiB
Python
import asyncio
|
|
import os
|
|
import tempfile
|
|
from asyncio import AbstractEventLoop
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
from shutil import rmtree
|
|
from threading import Thread
|
|
from unittest import IsolatedAsyncioTestCase
|
|
from unittest.mock import AsyncMock, MagicMock, mock_open, patch
|
|
|
|
from aio_pika import IncomingMessage
|
|
from aio_pika.abc import AbstractRobustConnection
|
|
|
|
from amqp.adapter import file_handler
|
|
from amqp.adapter.file_handler import StreamingFileHandler
|
|
|
|
|
|
class TestFileHandler(IsolatedAsyncioTestCase):
|
|
def test_ensure_directory_exists(self):
|
|
"""
|
|
Test function for ensure_directory_exists. Creates a temporary directory
|
|
and files to test the versioning logic.
|
|
"""
|
|
|
|
# Create a temporary directory
|
|
file_handler = StreamingFileHandler()
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
# Test case 1: File does not exist
|
|
filepath1 = os.path.join(temp_dir, "test1.txt")
|
|
result1 = file_handler.ensure_directory_exists(filepath1)
|
|
assert result1 == filepath1, f"Test Case 1 Failed: {result1} != {filepath1}"
|
|
# Create the file
|
|
with open(filepath1, "w") as f:
|
|
f.write("test")
|
|
|
|
# Test case 2: File exists, should create versioned file
|
|
filepath2 = os.path.join(temp_dir, "test2.txt")
|
|
result2 = file_handler.ensure_directory_exists(filepath2)
|
|
assert result2 == filepath2, f"Test Case 2 Failed: {result2} != {filepath2}"
|
|
with open(filepath2, "w") as f:
|
|
f.write("test")
|
|
result2_v1 = file_handler.ensure_directory_exists(filepath2)
|
|
assert result2_v1 == os.path.join(
|
|
temp_dir, "test2.1.txt"
|
|
), f"Test Case 2 Version 1 Failed: {result2_v1} != {os.path.join(temp_dir, 'test2.1.txt')}"
|
|
|
|
# Create a few versioned files
|
|
with open(os.path.join(temp_dir, "test3.txt"), "w") as f:
|
|
f.write("test")
|
|
with open(os.path.join(temp_dir, "test3.1.txt"), "w") as f:
|
|
f.write("test")
|
|
with open(os.path.join(temp_dir, "test3.2.txt"), "w") as f:
|
|
f.write("test")
|
|
filepath3 = os.path.join(temp_dir, "test3.txt")
|
|
result3_v3 = file_handler.ensure_directory_exists(filepath3)
|
|
assert result3_v3 == os.path.join(
|
|
temp_dir, "test3.3.txt"
|
|
), f"Test Case 3 Version 3 Failed: {result3_v3} != {os.path.join(temp_dir, 'test3.3.txt')}"
|
|
|
|
print("All test cases passed!")
|
|
|
|
async def test_download_from_queue_to_file(self):
|
|
loop = asyncio.new_event_loop()
|
|
tempFile = tempfile.mktemp()
|
|
file_handler = StreamingFileHandler()
|
|
with patch.object(
|
|
StreamingFileHandler, "_get_connection", new_callable=AsyncMock
|
|
) as mock_connect:
|
|
mock_connection: AbstractRobustConnection = MagicMock(spec=AbstractRobustConnection)
|
|
mock_connect.return_value = mock_connection
|
|
mock_connection.channel = AsyncMock()
|
|
mock_channel = mock_connection.channel.return_value
|
|
mock_channel.get_queue = AsyncMock()
|
|
mock_queue = mock_channel.get_queue.return_value
|
|
mock_queue.iterator = MagicMock()
|
|
mock_queue_iter = mock_queue.iterator().__aenter__.return_value
|
|
mock_message: IncomingMessage = MagicMock(spec=IncomingMessage)
|
|
mock_message.body = b"test_data"
|
|
mock_message.body_size = len(b"test_data")
|
|
mock_message.headers = {"eof": 1}
|
|
mock_queue_iter.__aiter__.return_value = [mock_message]
|
|
|
|
# Act
|
|
result = await file_handler._download_from_queue_to_file(
|
|
"test_rabbitmq",
|
|
"test_queue",
|
|
filepath=tempFile,
|
|
loop=loop,
|
|
)
|
|
# clean up
|
|
os.remove(tempFile)
|
|
|
|
# Assert
|
|
assert result == (len(b"test_data"), 1)
|
|
mock_connect.assert_called_once_with("test_rabbitmq", loop=loop)
|
|
mock_channel.get_queue.assert_called_once_with("test_queue", ensure=False)
|
|
|
|
async def test_download_file_success(self):
|
|
# Arrange
|
|
file_handler = StreamingFileHandler()
|
|
loop = asyncio.new_event_loop()
|
|
loop_thread = Thread(target=loop.run_forever, daemon=True)
|
|
loop_thread.start()
|
|
rabbitmq_url = "amqp://test.rabbitmq"
|
|
file_info = {
|
|
"filename": "test_file.txt",
|
|
"size": 1024,
|
|
"streamName": "test_stream",
|
|
"queue_name": "test_queue",
|
|
"key": "test_key",
|
|
}
|
|
message_data = {}
|
|
output_dir = tempfile.mkdtemp()
|
|
|
|
file_handler.ensure_directory_exists = MagicMock(return_value=f"{output_dir}/test_file.txt")
|
|
file_handler._download_from_queue_to_file = AsyncMock(return_value=(123, 1))
|
|
|
|
# Act
|
|
result = await file_handler.download_file(
|
|
loop, rabbitmq_url, file_info, message_data, output_dir
|
|
)
|
|
|
|
# clean up
|
|
loop.stop()
|
|
rmtree(output_dir, ignore_errors=True)
|
|
|
|
# Assert
|
|
assert result == (123, 1)
|
|
file_handler.ensure_directory_exists.assert_called_once_with(f"{output_dir}/test_file.txt")
|
|
file_handler._download_from_queue_to_file.assert_called_once_with(
|
|
"amqp://test.rabbitmq", "test_queue", f"{output_dir}/test_file.txt", loop
|
|
)
|
|
|
|
async def test_download_file_queue_not_found(self):
|
|
# Arrange
|
|
file_handler = StreamingFileHandler()
|
|
loop = MagicMock(spec=AbstractEventLoop)
|
|
rabbitmq_url = "amqp://test.rabbitmq"
|
|
file_info = {
|
|
"filename": "test_file.txt",
|
|
"size": 1024,
|
|
"streamName": "test_stream",
|
|
"queue_name": None,
|
|
"key": "test_key",
|
|
}
|
|
message_data = {}
|
|
output_dir = "test_output"
|
|
|
|
# Act
|
|
result = await file_handler.download_file(
|
|
loop, rabbitmq_url, file_info, message_data, output_dir
|
|
)
|
|
|
|
# Assert
|
|
assert result == (0, 2)
|
|
|
|
async def test_download_from_queue_to_buffer(self):
|
|
loop = asyncio.new_event_loop()
|
|
file_handler = StreamingFileHandler()
|
|
with patch.object(
|
|
StreamingFileHandler, "_get_connection", new_callable=AsyncMock
|
|
) as mock_connect:
|
|
mock_connection: AbstractRobustConnection = MagicMock(spec=AbstractRobustConnection)
|
|
mock_connect.return_value = mock_connection
|
|
mock_connection.channel = AsyncMock()
|
|
mock_channel = mock_connection.channel.return_value
|
|
mock_channel.get_queue = AsyncMock()
|
|
mock_queue = mock_channel.get_queue.return_value
|
|
mock_queue.iterator = MagicMock()
|
|
mock_queue_iter = mock_queue.iterator().__aenter__.return_value
|
|
mock_message: IncomingMessage = MagicMock(spec=IncomingMessage)
|
|
mock_message.body = b"test_data"
|
|
mock_message.body_size = len(b"test_data")
|
|
mock_message.headers = {"eof": 1}
|
|
mock_queue_iter.__aiter__.return_value = [mock_message]
|
|
|
|
# Act
|
|
result = await file_handler._download_from_queue_to_buffer(
|
|
"test_rabbitmq",
|
|
"test_queue",
|
|
loop=loop,
|
|
)
|
|
|
|
# Assert
|
|
assert result == (b"test_data", 1)
|
|
mock_connect.assert_called_once_with("test_rabbitmq", loop=loop)
|
|
mock_channel.get_queue.assert_called_once_with("test_queue", ensure=False)
|
|
|
|
async def test_download_buffer_success(self):
|
|
# Arrange
|
|
file_handler = StreamingFileHandler()
|
|
loop = asyncio.new_event_loop()
|
|
loop_thread = Thread(target=loop.run_forever, daemon=True)
|
|
loop_thread.start()
|
|
rabbitmq_url = "amqp://test.rabbitmq"
|
|
|
|
file_handler._download_from_queue_to_buffer = AsyncMock(return_value=(b"hello", 1))
|
|
|
|
# Act
|
|
result = await file_handler.download_buffer(loop, rabbitmq_url, "test_queue")
|
|
|
|
# clean up
|
|
loop.stop()
|
|
|
|
# Assert
|
|
assert result == (b"hello", 1)
|
|
file_handler._download_from_queue_to_buffer.assert_called_once_with(
|
|
"amqp://test.rabbitmq", "test_queue", loop
|
|
)
|
|
|
|
async def test_download_buffer_queue_not_found(self):
|
|
# Arrange
|
|
file_handler = StreamingFileHandler()
|
|
loop = MagicMock(spec=AbstractEventLoop)
|
|
rabbitmq_url = "amqp://test.rabbitmq"
|
|
|
|
# Act
|
|
result = await file_handler.download_buffer(loop, rabbitmq_url, None)
|
|
|
|
# Assert
|
|
assert result == (0, 2)
|
|
|
|
async def test_on_message_no_files_to_download(self):
|
|
handler = StreamingFileHandler()
|
|
message = MagicMock()
|
|
message.body = b"{}"
|
|
json_data = b'{"files": {}}'
|
|
rabbitmq_url = "amqp://localhost"
|
|
loop = AsyncMock()
|
|
output_dir = "downloaded_files"
|
|
|
|
handler.download_file = AsyncMock()
|
|
handler.download_file.return_value = (123, 1)
|
|
with patch("amqp.adapter.file_handler.logging_info"):
|
|
result = await handler.on_message(message, json_data, rabbitmq_url, loop, output_dir)
|
|
|
|
assert isinstance(result, defaultdict)
|
|
self.assertEqual(result, defaultdict())
|
|
|
|
async def test_on_message_with_files_to_download(self):
|
|
handler = StreamingFileHandler()
|
|
message = MagicMock()
|
|
message.body = b'{"fileStream1": "queue1"}'
|
|
json_data = b'{"files": {"key1": [{"streamName": "fileStream1"}]}}'
|
|
rabbitmq_url = "amqp://localhost"
|
|
loop = AsyncMock()
|
|
output_dir = "downloaded_files"
|
|
|
|
handler.download_file = AsyncMock()
|
|
|
|
with patch("amqp.adapter.file_handler.logging_info") as mock_logging_info:
|
|
result = await handler.on_message(message, json_data, rabbitmq_url, loop, output_dir)
|
|
|
|
self.assertEqual(result, file_handler.download_locations)
|
|
handler.download_file.assert_awaited_once()
|
|
mock_logging_info.assert_any_call(
|
|
f" [{message.delivery_tag}] about to create task for {{'streamName': 'fileStream1', 'queue_name': 'queue1', 'key': 'key1'}}"
|
|
)
|
|
mock_logging_info.assert_called_with(f"All files downloaded. d-tag: {message.delivery_tag}")
|
|
|
|
async def test_on_message_handles_exception(self):
|
|
handler = StreamingFileHandler()
|
|
message = MagicMock()
|
|
message.body = b"invalid_json"
|
|
message.ack = AsyncMock()
|
|
json_data = b"invalid_json"
|
|
rabbitmq_url = "amqp://localhost"
|
|
loop = AsyncMock()
|
|
output_dir = "downloaded_files"
|
|
|
|
with patch("amqp.adapter.file_handler.logging_error") as mock_logging_error:
|
|
await handler.on_message(message, json_data, rabbitmq_url, loop, output_dir)
|
|
|
|
mock_logging_error.assert_called()
|
|
|
|
async def test_publish_file_file_not_found(self):
|
|
file_handler = StreamingFileHandler()
|
|
loop = asyncio.get_event_loop()
|
|
exchange_mock = AsyncMock()
|
|
file_path = Path("/path/to/nonexistent_file.txt")
|
|
|
|
with self.assertRaises(FileNotFoundError):
|
|
await file_handler.publish_file(
|
|
exchange=exchange_mock,
|
|
file_path=file_path,
|
|
routing_key="test_routing_key",
|
|
loop=loop,
|
|
)
|
|
|
|
async def test_publish_file_success(self):
|
|
file_handler = StreamingFileHandler()
|
|
loop = asyncio.get_event_loop()
|
|
exchange_mock = AsyncMock()
|
|
file_path = Path("/path/to/test_file.txt")
|
|
data = b"test-data"
|
|
with (
|
|
patch("amqp.adapter.file_handler.Path.is_file", return_value=True),
|
|
patch("amqp.adapter.file_handler.os.path.getsize", return_value=len(data)),
|
|
patch("builtins.open", new=mock_open(read_data=data)) as mock_file,
|
|
):
|
|
result = await file_handler.publish_file(
|
|
exchange=exchange_mock,
|
|
file_path=file_path,
|
|
routing_key="test_routing_key",
|
|
loop=loop,
|
|
max_chunk_size=512,
|
|
)
|
|
|
|
mock_file.assert_called_with(file_path, "rb")
|
|
self.assertEqual(result, ("test_routing_key", len(data)))
|