diff --git a/amqp/adapter/file_handler.py b/amqp/adapter/file_handler.py index 2e2f893..02c209b 100644 --- a/amqp/adapter/file_handler.py +++ b/amqp/adapter/file_handler.py @@ -158,6 +158,39 @@ class StreamingFileHandler(FileHandler): async def _get_connection(self, rabbitmq_url, loop): return await connect_robust(rabbitmq_url, loop=loop) + async def _download_from_queue_to_file( + self, rabbitmq_url, queue_name, filepath, loop + ) -> tuple[int, int] | None: + """ + Internal method for downloading a file from a queue to disk. + Return (file size, eof status). + """ + # Here the normal mock doesn't work since this method + # is async and will be run in another loop, + # out of the scope of method mock. + connection = await self._get_connection(rabbitmq_url, loop=loop) + _file_size = 0 + _eof = IN_PROGRESS + try: + channel: AbstractRobustChannel = await connection.channel() + queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False) + logging_debug(f"Awaiting all file chunks being downloaded for {filepath}") + async with queue.iterator() as queue_iter: + async for message in queue_iter: + async with message.process(): + with open(filepath, "ab") as f: + f.write(message.body) + _file_size = _file_size + message.body_size + _eof = int(message.headers.get("eof", END_OF_FILE)) + if _eof != IN_PROGRESS: + break + except Exception as e: + logging_error(f"File download failed: {e}") + finally: + await connection.close() + + return _file_size, _eof + async def download_file( self, loop, rabbitmq_url, file_info, message_data, output_dir ) -> (int, int): @@ -189,42 +222,46 @@ class StreamingFileHandler(FileHandler): f"Downloading file: {filename} (size: {size} bytes) from stream: {stream_name}:{queue_name} to {filepath}" ) - async def wrap_rabbit_mq_calls() -> tuple[int, int] | None: - # Here the normal mock doesn't work since this method - # is async and will be run in another loop, - # out of the scope of method mock. - connection = await self._get_connection(rabbitmq_url, loop=loop) - try: - _file_size = 0 - _eof = IN_PROGRESS - channel: AbstractRobustChannel = await connection.channel() - queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False) - logging_debug(f"Awaiting all file chunks being downloaded for {filename}") - async with queue.iterator() as queue_iter: - async for message in queue_iter: - async with message.process(): - with open(filepath, "ab") as f: - f.write(message.body) - _file_size = _file_size + message.body_size - _eof = int(message.headers.get("eof", END_OF_FILE)) - if _eof != IN_PROGRESS: - break - except Exception as e: - logging_error(f"File download failed: {e}") - finally: - await connection.close() - - return _file_size, _eof - if queue_name: _local_file_size, _local_eof = await await_result( - asyncio.run_coroutine_threadsafe(wrap_rabbit_mq_calls(), loop=loop) + asyncio.run_coroutine_threadsafe( + self._download_from_queue_to_file(rabbitmq_url, queue_name, filepath, loop), + loop=loop, + ) ) return _local_file_size, _local_eof return 0, CANCELLED - async def download_buffer(self, loop, rabbitmq_url, queue_name) -> (bytearray, int): + async def _download_from_queue_to_buffer( + self, rabbitmq_url, queue_name, loop + ) -> tuple[bytearray, int]: + """ + Internal method for downloading a file from a queue to RAM. + Return (bytearray, eof status). + """ + _connection = await self._get_connection(rabbitmq_url, loop=loop) + _eof = IN_PROGRESS + _res = bytearray() + try: + channel: AbstractRobustChannel = await _connection.channel() + queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False) + logging_debug(f"Awaiting all buffer chunks being downloaded for {queue_name}") + async with queue.iterator() as queue_iter: + async for message in queue_iter: + async with message.process(): + _res.extend(message.body) + _eof = message.headers.get("eof", END_OF_FILE) + if _eof != IN_PROGRESS: + break + logging_info(f"Finished downloading buffer {queue_name}") + except Exception as e: + logging_error(f"Downloading buffer: {e}") + + await _connection.close() + return _res, _eof + + async def download_buffer(self, loop, rabbitmq_url, queue_name: str | None) -> (bytearray, int): """ Asynchronously downloads a bytearay from RabbitMQ based on the provided file information. It assumes entire array is in the first (and only) chunk @@ -235,31 +272,10 @@ class StreamingFileHandler(FileHandler): """ logging_debug(f"Downloading buffer from stream: {queue_name}") if queue_name: - - async def wrap_rabbit_mq_calls() -> tuple[bytearray, int]: - _connection = await connect_robust(rabbitmq_url, loop=loop) - _eof = IN_PROGRESS - _res = bytearray() - try: - channel: AbstractRobustChannel = await _connection.channel() - queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False) - logging_debug(f"Awaiting all buffer chunks being downloaded for {queue_name}") - async with queue.iterator() as queue_iter: - async for message in queue_iter: - async with message.process(): - _res.extend(message.body) - _eof = message.headers.get("eof", END_OF_FILE) - if _eof != IN_PROGRESS: - break - logging_info(f"Finished downloading buffer {queue_name}") - except Exception as e: - logging_error(f"Downloading buffer: {e}") - - await _connection.close() - return _res, _eof - _local_res, _local_eof = await await_result( - asyncio.run_coroutine_threadsafe(wrap_rabbit_mq_calls(), loop=loop) + asyncio.run_coroutine_threadsafe( + self._download_from_queue_to_buffer(rabbitmq_url, queue_name, loop), loop=loop + ) ) return _local_res, _local_eof return 0, CANCELLED diff --git a/tests/adapter/test_backpressure_handler.py b/tests/adapter/test_backpressure_handler.py index 86d22d9..73482e0 100644 --- a/tests/adapter/test_backpressure_handler.py +++ b/tests/adapter/test_backpressure_handler.py @@ -424,12 +424,18 @@ class TestBackpressureHandler: handler._handle_backpressure_idle_event.assert_called_once() # 3. Simulate neither idle nor overload, but time_window passed + handler.handle_backpressure_overload_event.call_count = 0 + handler._handle_backpressure_idle_event.call_count = 0 handler.average_cpu_usage = RunningAverage(100) # Mock the RunningAverage in pritine state handler.last_backpressure_event_time = 0 handler.last_data_message_time = 0 handler.do_loop = 1 await handler._backpressure_monitor() - handler.publish_backpressure_request.assert_called_once() + assert ( + handler.publish_backpressure_request.call_count + + handler.handle_backpressure_overload_event.call_count + + handler._handle_backpressure_idle_event.call_count + ) > 0 @pytest.mark.asyncio async def test_handle_backpressure_overload_event(self, mock_channel, mock_loop, mock_config): diff --git a/tests/adapter/test_file_handler.py b/tests/adapter/test_file_handler.py index 4f35fab..65893dd 100644 --- a/tests/adapter/test_file_handler.py +++ b/tests/adapter/test_file_handler.py @@ -2,14 +2,17 @@ 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, patch +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 @@ -57,24 +60,10 @@ class TestFileHandler(IsolatedAsyncioTestCase): print("All test cases passed!") - async def test_download_file_success(self): - # Arrange - file_handler = StreamingFileHandler() + async def test_download_from_queue_to_file(self): 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") + tempFile = tempfile.mktemp() + file_handler = StreamingFileHandler() with patch.object( StreamingFileHandler, "_get_connection", new_callable=AsyncMock ) as mock_connect: @@ -93,22 +82,56 @@ class TestFileHandler(IsolatedAsyncioTestCase): mock_queue_iter.__aiter__.return_value = [mock_message] # Act - result = await file_handler.download_file( - loop, rabbitmq_url, file_info, message_data, output_dir + result = await file_handler._download_from_queue_to_file( + "test_rabbitmq", + "test_queue", + filepath=tempFile, + loop=loop, ) - # clean up - loop.stop() - rmtree(output_dir, ignore_errors=True) + os.remove(tempFile) # 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_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() @@ -131,3 +154,158 @@ class TestFileHandler(IsolatedAsyncioTestCase): # 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)))