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

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