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)))