Add unit tests for asyncio download methods

STONEBLD-1356

Signed-off-by: Felipe de Almeida <[email protected]>

pyproject.toml

@@ -50,6 +50,7 @@ test = [
   "GitPython",
   "jsonschema",
   "pytest",
+  "pytest-asyncio",
   "pytest-cov",
   "pytest-env",
   "pyyaml",

requirements-extras.txt

@@ -627,8 +627,13 @@ pytest==7.4.0 \
     --hash=sha256:b4bf8c45bd59934ed84001ad51e11b4ee40d40a1229d2c79f9c592b0a3f6bd8a
     # via
     #   cachi2 (pyproject.toml)
+    #   pytest-asyncio
     #   pytest-cov
     #   pytest-env
+pytest-asyncio==0.21.0 \
+    --hash=sha256:2b38a496aef56f56b0e87557ec313e11e1ab9276fc3863f6a7be0f1d0e415e1b \
+    --hash=sha256:f2b3366b7cd501a4056858bd39349d5af19742aed2d81660b7998b6341c7eb9c
+    # via cachi2 (pyproject.toml)
 pytest-cov==4.1.0 \
     --hash=sha256:3904b13dfbfec47f003b8e77fd5b589cd11904a21ddf1ab38a64f204d6a10ef6 \
     --hash=sha256:6ba70b9e97e69fcc3fb45bfeab2d0a138fb65c4d0d6a41ef33983ad114be8c3a

tests/unit/conftest.py

@@ -1,6 +1,10 @@
+import asyncio
+import random
 import tarfile
 from pathlib import Path
+from unittest.mock import MagicMock
 
+import aiohttp_retry
 import pytest
 
 
@@ -17,3 +21,35 @@ def golang_repo_path(data_dir: Path, tmp_path: Path) -> Path:
         tar.extractall(tmp_path)
 
     return tmp_path / "golang_git_repo"
+
+
+@pytest.fixture
+def mock_async_download_binary_file() -> MagicMock:
+    async def mock_download_binary_file(
+        session: aiohttp_retry.RetryClient,
+        url: str,
+        download_path: str,
+    ) -> dict[str, str]:
+        # Simulate a file download by sleeping for a random duration
+        await asyncio.sleep(random.uniform(0.1, 0.5))
+
+        # Write some dummy data to the download path
+        with open(download_path, "wb") as file:
+            file.write(b"Mock file content")
+
+        # Return a dummy response indicating success
+        return {"status": "success", "url": url, "download_path": download_path}
+
+    return MagicMock(side_effect=mock_download_binary_file)
+
+
+class MockReadChunk:
+    def __init__(self) -> None:
+        """Create a call count."""
+        self.call_count = 0
+
+    async def read_chunk(self, size: int) -> bytes:
+        """Return a non-empty chunk for the first and second call, then an empty chunk."""
+        self.call_count += 1
+        chunks = {1: b"first_chunk-", 2: b"second_chunk-"}
+        return chunks.get(self.call_count, b"")

tests/unit/package_managers/test_general.py

@@ -1,7 +1,10 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
+import re
+from os import PathLike
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Dict, Optional, Union
 from unittest import mock
+from unittest.mock import MagicMock
 
 import pytest
 import requests
@@ -10,7 +13,13 @@
 from cachi2.core.config import get_config
 from cachi2.core.errors import FetchError
 from cachi2.core.package_managers import general
-from cachi2.core.package_managers.general import download_binary_file, pkg_requests_session
+from cachi2.core.package_managers.general import (
+    _async_download_binary_file,
+    async_download_files,
+    download_binary_file,
+    pkg_requests_session,
+)
+from tests.unit.conftest import MockReadChunk
 
 GIT_REF = "9a557920b2a6d4110f838506120904a6fda421a2"
 
@@ -132,3 +141,84 @@ def test_extract_git_info(url: str, nonstandard_info: Any) -> None:
     }
     info.update(nonstandard_info or {})
     assert general.extract_git_info(url) == info
+
+
+@pytest.mark.asyncio
+async def test_async_download_binary_file(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
+    url = "http://example.com/file.tar"
+    download_path = tmp_path / "file.tar"
+    chunk_size = 8192
+
+    # Create mock response and session
+    response, session = MagicMock(), MagicMock()
+    response.content.read = MockReadChunk().read_chunk
+
+    async def mock_aenter() -> MagicMock:
+        return response
+
+    session.get().__aenter__.side_effect = mock_aenter
+
+    await _async_download_binary_file(session, url, download_path, chunk_size=chunk_size)
+
+    # Verify that the file was downloaded correctly
+    with open(download_path, "rb") as f:
+        assert f.read() == b"first_chunk-second_chunk-"
+
+    # # Verify that session.get was called with the correct arguments
+    assert session.get.called
+    assert session.get.call_args == mock.call(url, auth=None, raise_for_status=True)
+
+    # Test unsuccessful download with an exception
+    exception_message = "This is a test exception message."
+    session.get().__aenter__.side_effect = Exception(exception_message)
+
+    with pytest.raises(Exception) as exc_info:
+        await _async_download_binary_file(session, url, download_path, chunk_size=chunk_size)
+
+    # Verify that the log message was correctly generated
+    assert f"Unsuccessful download: {url}" in caplog.text
+
+    # Verify that the exception was correctly captured and re-raised as FetchError
+    assert isinstance(exc_info.value, FetchError)
+    assert str(exc_info.value) == f"exception_name: Exception, details: {exception_message}"
+
+
+@pytest.mark.asyncio
+@mock.patch("cachi2.core.package_managers.general._async_download_binary_file")
+async def test_async_download_files(
+    mock_download_file: MagicMock,
+    tmp_path: Path,
+    mock_async_download_binary_file: MagicMock,
+) -> None:
+    files_to_download: Dict[str, Union[str, PathLike[str]]] = {
+        "file1": str(tmp_path / "path1"),
+        "file2": str(tmp_path / "path2"),
+        "file3": str(tmp_path / "path3"),
+    }
+
+    concurrency_limit = 2
+
+    # Set up the mock for _async_download_binary_file
+    mock_download_file.return_value = mock_async_download_binary_file
+
+    await async_download_files(files_to_download, concurrency_limit)
+
+    # Assert that mock_download_file was called for each file
+    assert mock_download_file.call_count == 3
+
+    # Assert that mock_download_file was called with the correct arguments
+    for call in mock_download_file.mock_calls:
+        # call looks like this:
+        # call(<AsyncMock name='RetryClient().__aenter__()' id='1'>, 'file1', '/tmp/path1')
+        file, path = re.findall(r"'([^']+)'", str(call))[-2:]
+        assert file, path in files_to_download.items()
+
+    # Set up the mock for _async_download_binary_file
+    exception_message = "This is a test exception message."
+    mock_download_file.side_effect = FetchError(exception_message)
+
+    with pytest.raises(FetchError) as exc_info:
+        await async_download_files(files_to_download, concurrency_limit)
+
+    # Verify that the exception was correctly captured and re-raised as FetchError
+    assert isinstance(exc_info.value, FetchError)