134 lines
5.3 KiB
Python
134 lines
5.3 KiB
Python
import asyncio
|
|
import os
|
|
import tempfile
|
|
from asyncio import AbstractEventLoop
|
|
from shutil import rmtree
|
|
from threading import Thread
|
|
from unittest import IsolatedAsyncioTestCase
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from aio_pika import IncomingMessage
|
|
from aio_pika.abc import AbstractRobustConnection
|
|
|
|
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_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")
|
|
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_file(
|
|
loop, rabbitmq_url, file_info, message_data, output_dir
|
|
)
|
|
|
|
# clean up
|
|
loop.stop()
|
|
rmtree(output_dir, ignore_errors=True)
|
|
|
|
# Assert
|
|
assert result == (len(b"test_data"), 1)
|
|
file_handler.ensure_directory_exists.assert_called_once_with(
|
|
f"{output_dir}/test_file.txt"
|
|
)
|
|
mock_connect.assert_called_once_with(rabbitmq_url, loop=loop)
|
|
mock_channel.get_queue.assert_called_once_with("test_queue", ensure=False)
|
|
|
|
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)
|