diff --git a/docs/task3_status.md b/docs/task3_status.md index b57f359..6de768e 100644 --- a/docs/task3_status.md +++ b/docs/task3_status.md @@ -128,7 +128,7 @@ The unit tests were written based on assumptions about the API that don't match - `src/core/utils/nfo_generator.py`: ✅ 19/19 tests passing - `src/core/utils/image_downloader.py`: ⚠️ 12/20 tests passing (8 require refactoring for context manager/attribute access) -- `src/core/services/tmdb_client.py`: ⚠️ 0/16 tests passing (async mocking challenges) +- `src/core/services/tmdb_client.py`: ⚠️ 0/16 tests passing (async mocking challenges) - `src/core/services/nfo_service.py`: ✅ 4 tests for update logic passing **Integration Tests:** diff --git a/src/core/utils/image_downloader.py b/src/core/utils/image_downloader.py index 80f0c61..0f52bfe 100644 --- a/src/core/utils/image_downloader.py +++ b/src/core/utils/image_downloader.py @@ -27,17 +27,25 @@ class ImageDownloadError(Exception): class ImageDownloader: """Utility for downloading and validating images. + Supports async context manager protocol for proper resource cleanup. + Attributes: max_retries: Maximum retry attempts for downloads timeout: Request timeout in seconds min_file_size: Minimum valid file size in bytes + session: Optional aiohttp session (managed internally) + + Example: + >>> async with ImageDownloader() as downloader: + ... await downloader.download_poster(url, path) """ def __init__( self, max_retries: int = 3, timeout: int = 60, - min_file_size: int = 1024 # 1 KB + min_file_size: int = 1024, # 1 KB + retry_delay: float = 1.0 ): """Initialize image downloader. @@ -45,10 +53,28 @@ class ImageDownloader: max_retries: Maximum retry attempts timeout: Request timeout in seconds min_file_size: Minimum valid file size in bytes + retry_delay: Delay between retries in seconds """ self.max_retries = max_retries self.timeout = timeout self.min_file_size = min_file_size + self.retry_delay = retry_delay + self.session: Optional[aiohttp.ClientSession] = None + + async def __aenter__(self): + """Enter async context manager.""" + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + """Exit async context manager and cleanup resources.""" + await self.close() + return False + + async def close(self): + """Close aiohttp session if open.""" + if self.session and not self.session.closed: + await self.session.close() + self.session = None async def download_image( self,