test(General): finished unit test for file handler, fixed unit test for backpressure handler

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

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

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):
+    async def test_download_from_queue_to_file(self):
-        # Arrange
-        file_handler = StreamingFileHandler()
         loop = asyncio.new_event_loop()
-        loop_thread = Thread(target=loop.run_forever, daemon=True)
+        tempFile = tempfile.mktemp()
-        loop_thread.start()
+        file_handler = StreamingFileHandler()
-        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:
@@ -93,22 +82,56 @@ class TestFileHandler(IsolatedAsyncioTestCase):
             mock_queue_iter.__aiter__.return_value = [mock_message]
 
             # Act
-            result = await file_handler.download_file(
+            result = await file_handler._download_from_queue_to_file(
-                loop, rabbitmq_url, file_info, message_data, output_dir
+                "test_rabbitmq",
+                "test_queue",
+                filepath=tempFile,
+                loop=loop,
             )
-
             # clean up
-            loop.stop()
+            os.remove(tempFile)
-            rmtree(output_dir, ignore_errors=True)
 
             # Assert
             assert result == (len(b"test_data"), 1)
-            file_handler.ensure_directory_exists.assert_called_once_with(
+            mock_connect.assert_called_once_with("test_rabbitmq", loop=loop)
-                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_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)))