diff --git a/amqp/adapter/file_downloader.py b/amqp/adapter/file_downloader.py index 83acfb8..ff078c6 100644 --- a/amqp/adapter/file_downloader.py +++ b/amqp/adapter/file_downloader.py @@ -15,11 +15,34 @@ END_OF_FILE = 1 CANCELLED = 2 -def ensure_directory_exists(filepath): - """Ensures that the directory for the given filepath exists.""" - dir_name = os.path.dirname(filepath) - if dir_name and not os.path.exists(dir_name): - os.makedirs(dir_name) +def ensure_directory_exists(filepath) -> str: + """ + Ensures that the directory for the given filepath exists, and handles file versioning. + + Args: + filepath: The path to the file. + + Returns: + The original filepath if the file does not exist, or a modified filepath + pointing to the next available version of the file if it does. + """ + # 1. Ensure directory exists + directory = os.path.dirname(filepath) + if directory and not os.path.exists(directory): + os.makedirs(directory) + + # 2. Handle file versioning + if not os.path.exists(filepath): + return filepath # File doesn't exist, return original path + + # File exists, find next available version + base, ext = os.path.splitext(filepath) + version = 1 + while True: + new_filepath = f"{base}.{version}{ext}" + if not os.path.exists(new_filepath): + return new_filepath # Found a free version, return this path + version += 1 async def download_file(loop, rabbitmq_url, file_info, message_data, output_dir) -> (int, int): @@ -39,8 +62,11 @@ async def download_file(loop, rabbitmq_url, file_info, message_data, output_dir) stream_name = file_info['streamName'] queue_name = file_info['queue_name'] filepath = os.path.join(output_dir, filename) - ensure_directory_exists(filepath) + filepath = ensure_directory_exists(filepath) # if filename exists, create a new (next) version download_locations[file_info['key']] = filepath + filename = os.path.basename(filepath) # if a new version was created, ensure to use that new version + file_info['filename'] = filename + eof = IN_PROGRESS fsize = 0 logging.info(f"Downloading file: {filename} (size: {size} bytes) from stream: {stream_name}:{queue_name} to {filepath}") @@ -67,7 +93,7 @@ async def download_file(loop, rabbitmq_url, file_info, message_data, output_dir) return (fsize, eof) -async def download_buffer(loop, rabbitmq_url, queue_name, ttl = 604800000) -> (bytearray, int): +async def download_buffer(loop, rabbitmq_url, queue_name) -> (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 @@ -97,7 +123,9 @@ async def download_buffer(loop, rabbitmq_url, queue_name, ttl = 604800000) -> (b logging.error(f"Error downloading buffer: {e}") tb = traceback.extract_tb(sys.exc_info()[2])[-1] logging.error(f" [E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {e}") - await connection.close() + + await connection.close() + return (res, eof) diff --git a/amqp/adapter/test_file_downloader.py b/amqp/adapter/test_file_downloader.py new file mode 100644 index 0000000..de24c72 --- /dev/null +++ b/amqp/adapter/test_file_downloader.py @@ -0,0 +1,46 @@ +import os +import tempfile +from unittest import TestCase + +import pytest + +from amqp.adapter.file_downloader import ensure_directory_exists + +class TestFileDownloader(TestCase): + 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 + with tempfile.TemporaryDirectory() as temp_dir: + # Test case 1: File does not exist + filepath1 = os.path.join(temp_dir, "test1.txt") + result1 = 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 = 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 = 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 = 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!") diff --git a/requirements.txt b/requirements.txt index 9f24c96..e97ce46 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,3 +7,4 @@ opentelemetry-exporter-otlp opentelemetry-exporter-prometheus opentelemetry-instrumentation-pika python_multipart +pytest