From 98fc52fc26544e82ef094d959ed0800ebac436fe Mon Sep 17 00:00:00 2001
From: Rui Hu <rui.hu@cleverthis.com>
Date: Tue, 13 May 2025 18:00:01 +0800
Subject: [PATCH] Partially finished the unit test for file handler

---
 amqp/adapter/file_handler.py       | 10 +++-
 tests/adapter/test_file_handler.py | 87 +++++++++++++++++++++++++++++-
 2 files changed, 94 insertions(+), 3 deletions(-)

diff --git a/amqp/adapter/file_handler.py b/amqp/adapter/file_handler.py
index 7c3e0bf..2e2f893 100644
--- a/amqp/adapter/file_handler.py
+++ b/amqp/adapter/file_handler.py
@@ -153,6 +153,11 @@ class StreamingFileHandler(FileHandler):
     def __init__(self):
         super().__init__()
 
+    # This method MUST NOT be static.
+    # For testing purpose, we will mock this method and inject a mock connection.
+    async def _get_connection(self, rabbitmq_url, loop):
+        return await connect_robust(rabbitmq_url, loop=loop)
+
     async def download_file(
         self, loop, rabbitmq_url, file_info, message_data, output_dir
     ) -> (int, int):
@@ -185,7 +190,10 @@ class StreamingFileHandler(FileHandler):
         )
 
         async def wrap_rabbit_mq_calls() -> tuple[int, int] | None:
-            connection = await connect_robust(rabbitmq_url, loop=loop)
+            # 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
diff --git a/tests/adapter/test_file_handler.py b/tests/adapter/test_file_handler.py
index 22b369c..4f35fab 100644
--- a/tests/adapter/test_file_handler.py
+++ b/tests/adapter/test_file_handler.py
@@ -1,11 +1,19 @@
+import asyncio
 import os
 import tempfile
-from unittest import TestCase
+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(TestCase):
+class TestFileHandler(IsolatedAsyncioTestCase):
     def test_ensure_directory_exists(self):
         """
         Test function for ensure_directory_exists.  Creates a temporary directory
@@ -48,3 +56,78 @@ class TestFileHandler(TestCase):
             ), 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)