xref: /aosp_15_r20/development/python-packages/fetchartifact/tests/test_fetchartifact.py (revision 90c8c64db3049935a07c6143d7fd006e26f8ecca) 
1#
2# Copyright (C) 2023 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16"""Tests for fetchartifact."""
17from typing import cast
18
19import pytest
20from aiohttp import ClientResponseError, ClientSession
21from aiohttp.test_utils import TestClient
22from aiohttp.web import Application, Request, Response
23
24from fetchartifact import ArtifactDownloader, fetch_artifact, fetch_artifact_chunked
25
26TEST_BUILD_ID = "1234"
27TEST_TARGET = "linux"
28TEST_ARTIFACT_NAME = "output.zip"
29TEST_DOWNLOAD_URL = (
30    f"/android/internal/build/v3/builds/{TEST_BUILD_ID}/{TEST_TARGET}/"
31    f"attempts/latest/artifacts/{TEST_ARTIFACT_NAME}/url"
32)
33TEST_RESPONSE = b"Hello, world!"
34
35
36@pytest.fixture(name="android_ci_client")
37async def fixture_android_ci_client(aiohttp_client: type[TestClient]) -> TestClient:
38    """Fixture for mocking the Android CI APIs."""
39
40    async def download(_request: Request) -> Response:
41        return Response(text=TEST_RESPONSE.decode("utf-8"))
42
43    app = Application()
44    app.router.add_get(TEST_DOWNLOAD_URL, download)
45    return await aiohttp_client(app)  # type: ignore
46
47
48async def test_fetch_artifact(android_ci_client: TestClient) -> None:
49    """Tests that the download URL is queried."""
50    assert TEST_RESPONSE == await fetch_artifact(
51        TEST_TARGET,
52        TEST_BUILD_ID,
53        TEST_ARTIFACT_NAME,
54        cast(ClientSession, android_ci_client),
55        query_url_base="",
56    )
57
58
59async def test_fetch_artifact_chunked(android_ci_client: TestClient) -> None:
60    """Tests that the full file contents are downloaded."""
61    assert [c.encode("utf-8") for c in TEST_RESPONSE.decode("utf-8")] == [
62        chunk
63        async for chunk in fetch_artifact_chunked(
64            TEST_TARGET,
65            TEST_BUILD_ID,
66            TEST_ARTIFACT_NAME,
67            cast(ClientSession, android_ci_client),
68            chunk_size=1,
69            query_url_base="",
70        )
71    ]
72
73
74async def test_failure_raises(android_ci_client: TestClient) -> None:
75    """Tests that fetch failure raises an exception."""
76    with pytest.raises(ClientResponseError):
77        await fetch_artifact(
78            TEST_TARGET,
79            TEST_BUILD_ID,
80            TEST_ARTIFACT_NAME,
81            cast(ClientSession, android_ci_client),
82            query_url_base="/bad",
83        )
84
85    with pytest.raises(ClientResponseError):
86        async for _chunk in fetch_artifact_chunked(
87            TEST_TARGET,
88            TEST_BUILD_ID,
89            TEST_ARTIFACT_NAME,
90            cast(ClientSession, android_ci_client),
91            query_url_base="/bad",
92        ):
93            pass
94
95
96class TestDownloader(ArtifactDownloader):
97    """Downloader which tracks calls to on_artifact_size and after_chunk."""
98
99    def __init__(self, target: str, build_id: str, artifact_name: str) -> None:
100        super().__init__(target, build_id, artifact_name, query_url_base="")
101        self.reported_content_length: int | None = None
102        self.reported_chunk_sizes: list[int] = []
103
104    def on_artifact_size(self, size: int) -> None:
105        super().on_artifact_size(size)
106        assert self.reported_content_length is None
107        self.reported_content_length = size
108
109    def after_chunk(self, size: int) -> None:
110        super().after_chunk(size)
111        self.reported_chunk_sizes.append(size)
112
113
114async def test_downloader_progress_reports(android_ci_client: TestClient) -> None:
115    """Tests that progress is reported when using ArtifactDownloader."""
116    downloader = TestDownloader(TEST_TARGET, TEST_BUILD_ID, TEST_ARTIFACT_NAME)
117
118    assert [b"Hell", b"o, w", b"orld", b"!"] == [
119        chunk
120        async for chunk in downloader.download(
121            cast(ClientSession, android_ci_client), chunk_size=4
122        )
123    ]
124    assert downloader.reported_content_length == len(TEST_RESPONSE.decode("utf-8"))
125    assert downloader.reported_chunk_sizes == [4, 4, 4, 1]
126
127
128@pytest.mark.requires_network
129async def test_real_artifact() -> None:
130    """Tests with a real artifact. Requires an internet connection."""
131    async with ClientSession() as session:
132        contents = await fetch_artifact("linux", "9945621", "logs/SUCCEEDED", session)
133        assert contents == b"1681499053\n"
134