# Copyright 2022-2023 The Debusine developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.

"""Unit tests for the Debusine class."""

import contextlib
import datetime
import hashlib
import io
import json
import logging
import os
import stat
import tarfile
import textwrap
import urllib
from collections.abc import Callable
from pathlib import Path
from typing import Any
from unittest import TestCase, mock
from unittest.mock import MagicMock, call
from urllib.parse import urlencode

import aiohttp

try:
    import pydantic.v1 as pydantic
except ImportError:
    import pydantic as pydantic  # type: ignore

import requests

import responses

from debusine.artifacts import WorkRequestDebugLogs
from debusine.artifacts.models import ArtifactCategory
from debusine.client.debusine import Debusine
from debusine.client.models import (
    ArtifactResponse,
    CollectionItemType,
    CreateWorkflowRequest,
    FileResponse,
    FilesResponseType,
    LookupMultipleResponse,
    LookupSingleResponse,
    RelationCreateRequest,
    RelationResponse,
    RemoteArtifact,
    WorkRequestRequest,
    WorkflowTemplateRequest,
    model_to_json_serializable_dict,
)
from debusine.client.tests import server
from debusine.client.tests.server import DebusineAioHTTPTestCase
from debusine.tasks.models import LookupMultiple
from debusine.test import TestHelpersMixin
from debusine.test.utils import (
    create_artifact_response,
    create_listing_response,
    create_work_request_response,
    create_workflow_template_response,
)


class DebusineTests(TestHelpersMixin, TestCase):
    """Tests for the Debusine class."""

    def setUp(self):
        """Initialize tests."""
        self.api_url = "https://debusine.debian.org/api"

        self.token = "token"
        self.logger = mock.create_autospec(spec=logging.Logger)
        self.debusine = Debusine(self.api_url, self.token, logger=self.logger)

    @responses.activate
    def test_work_request_iter(self):
        """Debusine.work_request_iter yields work requests."""
        work_requests = [
            create_work_request_response(id=1),
            create_work_request_response(id=2),
        ]
        listing_response = create_listing_response(*work_requests)

        responses.add(
            responses.GET,
            f'{self.debusine.api_url}/work-request/',
            json=model_to_json_serializable_dict(listing_response),
        )

        self.assertEqual(list(self.debusine.work_request_iter()), work_requests)
        self.assert_token_key_included_in_all_requests('token')

    @responses.activate
    def test_work_request_get(self):
        """Debusine.work_request_get returns the work request."""
        work_request_id = 10
        work_request_response = create_work_request_response(id=work_request_id)

        responses.add(
            responses.GET,
            f'{self.debusine.api_url}/work-request/{work_request_id}/',
            json=model_to_json_serializable_dict(work_request_response),
        )

        self.assertEqual(
            self.debusine.work_request_get(work_request_id),
            work_request_response,
        )
        self.assert_token_key_included_in_all_requests('token')

    @responses.activate
    def test_work_request_update(self) -> None:
        """Debusine.work_request_update patches the work request."""
        work_request_id = 10
        work_request_response = create_work_request_response(
            id=work_request_id, priority_adjustment=100
        )

        responses.add(
            responses.PATCH,
            f"{self.debusine.api_url}/work-request/{work_request_id}/",
            json=model_to_json_serializable_dict(work_request_response),
        )

        self.assertEqual(
            self.debusine.work_request_update(
                work_request_id, priority_adjustment=100
            ),
            work_request_response,
        )
        self.assert_token_key_included_in_all_requests("token")

    def verify_work_request_create_success(self, workspace: str | None = None):
        """
        Verify Debusine.work_request_create().

        :param workspace: if it's None: assert that the client's request
          does not contain "workspace" (it's assigned by the server). The
          server response contains workspace="System".

          If it's not None: the request contains "workspace" with the
          correct value.
        """
        expected_work_request = workspace if workspace is not None else "System"

        # response always contains "workspace". It's "system" if it was
        # no in the request
        expected_work_request_response = create_work_request_response(
            workspace=expected_work_request
        )

        responses.add(
            responses.POST,
            f'{self.debusine.api_url}/work-request/',
            json=model_to_json_serializable_dict(
                expected_work_request_response
            ),
        )

        work_request_fields: dict[str, Any] = {
            "task_name": "sbuild",
            "task_data": {},
            "event_reactions": {},
            "workspace": workspace,
        }
        # these fields are always posted
        fields_to_post = {"task_name", "task_data", "event_reactions"}

        # "workspace" is posted only if it's not None
        if workspace is not None:
            fields_to_post.add("workspace")

        # work_request_to_post can contain "workspace=None". The method
        # Debusine.work_request_create will remove this key if it's None
        work_request_to_post = WorkRequestRequest(**work_request_fields)

        actual_work_request_response = self.debusine.work_request_create(
            work_request_to_post
        )
        self.assertEqual(
            actual_work_request_response, expected_work_request_response
        )

        self.assert_token_key_included_in_all_requests(self.token)

        # The posted (over the wire) WorkRequestRequest does not contain
        # "workspace" if it was None
        self.assertEqual(
            json.loads(responses.calls[0].request.body or "null"),
            work_request_to_post.dict(include=fields_to_post),
        )

    @responses.activate
    def test_work_request_create_success(self):
        """
        Post a new work request for workspace "Testing".

        The server will assign it to "Testing".
        """
        self.verify_work_request_create_success(workspace="Testing")

    def test_work_request_completed_path(self):
        """Debusine.work_request_completed_path() return expected path."""
        work_request_id = 67
        self.assertEqual(
            Debusine.work_request_completed_path(work_request_id),
            f"/work-request/{work_request_id}/completed/",
        )

    @responses.activate
    def test_work_request_completed_update(self):
        """work_request_completed_path() make the expected HTTP request."""
        work_request_id = 68

        responses.add(
            responses.PUT,
            f"{self.api_url}/1.0"
            + self.debusine.work_request_completed_path(work_request_id),
            "",
        )

        self.debusine.work_request_completed_update(work_request_id, "success")

    @responses.activate
    def test_work_request_create_no_workspace_success(self):
        """
        Post a new work request without a workspace.

        Assert that the request does not contain "workspace" (not even
        None, as the server expects no workspace field and assigns it to the
        default).
        """
        self.verify_work_request_create_success(workspace=None)

    def verify_workflow_template_create_success(
        self, workspace: str | None = None
    ):
        """
        Verify Debusine.workflow_template_create().

        :param workspace: if it's None: assert that the client's request
          does not contain "workspace" (it's assigned by the server). The
          server response contains workspace="System".

          If it's not None: the request contains "workspace" with the
          correct value.
        """
        # Response always contains "workspace".  It's "system" if it was not
        # in the request.
        expected_workflow_template_response = create_workflow_template_response(
            workspace=(workspace if workspace is not None else "System")
        )

        responses.add(
            responses.POST,
            f'{self.debusine.api_url}/workflow-template/',
            json=model_to_json_serializable_dict(
                expected_workflow_template_response
            ),
            status=requests.codes.created,
        )

        workflow_template_fields: dict[str, Any] = {
            "name": "test",
            "task_name": "noop",
            "task_data": {},
            "workspace": workspace,
            "priority": 0,
        }
        # these fields are always posted
        fields_to_post = {"name", "task_name", "task_data", "priority"}

        # "workspace" is posted only if it's not None
        if workspace is not None:
            fields_to_post.add("workspace")

        # workflow_template_to_post can contain "workspace=None". The method
        # Debusine.workflow_template_create will remove this key if it's
        # None.
        workflow_template_to_post = WorkflowTemplateRequest(
            **workflow_template_fields
        )

        actual_workflow_template_response = (
            self.debusine.workflow_template_create(workflow_template_to_post)
        )
        self.assertEqual(
            actual_workflow_template_response,
            expected_workflow_template_response,
        )

        self.assert_token_key_included_in_all_requests(self.token)

        # The posted (over the wire) WorkflowTemplateRequest does not
        # contain "workspace" if it was None.
        self.assertEqual(
            json.loads(responses.calls[0].request.body or "null"),
            workflow_template_to_post.dict(include=fields_to_post),
        )

    @responses.activate
    def test_workflow_template_create_success(self):
        """Post a new workflow template."""
        self.verify_workflow_template_create_success(workspace="Testing")

    @responses.activate
    def test_workflow_template_create_no_workspace_success(self):
        """Post a new workflow template."""
        self.verify_workflow_template_create_success(workspace=None)

    def verify_workflow_create_success(self, workspace: str | None = None):
        """
        Verify Debusine.workflow_create().

        :param workspace: if it's None: assert that the client's request
          does not contain "workspace" (it's assigned by the server). The
          server response contains workspace="System".

          If it's not None: the request contains "workspace" with the
          correct value.
        """
        # Response always contains "workspace".  It's "system" if it was not
        # in the request.
        expected_workflow_response = create_work_request_response(
            workspace=(workspace if workspace is not None else "System")
        )

        responses.add(
            responses.POST,
            f"{self.debusine.api_url}/workflow/",
            json=model_to_json_serializable_dict(expected_workflow_response),
            status=requests.codes.created,
        )

        workflow_fields: dict[str, Any] = {
            "template_name": "test",
            "task_data": {},
            "workspace": workspace,
        }
        # these fields are always posted
        fields_to_post = {"template_name", "task_data"}

        # "workspace" is posted only if it's not None
        if workspace is not None:
            fields_to_post.add("workspace")

        # workflow_to_post can contain "workspace=None". The method
        # Debusine.workflow_create will remove this key if it's None.
        workflow_to_post = CreateWorkflowRequest(**workflow_fields)

        actual_workflow_response = self.debusine.workflow_create(
            workflow_to_post
        )
        self.assertEqual(actual_workflow_response, expected_workflow_response)

        self.assert_token_key_included_in_all_requests(self.token)

        # The posted (over the wire) CreateWorkflowRequest does not contain
        # "workspace" if it was None.
        self.assertEqual(
            json.loads(responses.calls[0].request.body or "null"),
            workflow_to_post.dict(include=fields_to_post),
        )

    @responses.activate
    def test_workflow_create_success(self):
        """Post a new workflow."""
        self.verify_workflow_create_success(workspace="Testing")

    @responses.activate
    def test_workflow_create_no_workspace_success(self):
        """Post a new workflow with no workspace."""
        self.verify_workflow_create_success(workspace=None)

    def files_for_artifact(self) -> dict[str, Path]:
        """Return dictionary with files to upload in an artifact."""
        return {
            "README": self.create_temporary_file(),
            ".dirstamp": self.create_temporary_file(),
        }

    @responses.activate
    def test_artifact_get(self):
        """Test artifact_get() return ArtifactResponse."""
        artifact_id = 2
        workspace = "test"
        download_artifact_tar_gz_url = "https://example.com/path/"
        category = "Testing"
        created_at = datetime.datetime.utcnow()
        data = {}

        get_artifact_response = create_artifact_response(
            id=artifact_id,
            workspace=workspace,
            category=category,
            created_at=created_at,
            data=data,
            download_tar_gz_url=download_artifact_tar_gz_url,
        )

        responses.add(
            responses.GET,
            f"{self.api_url}/1.0/artifact/{get_artifact_response.id}/",
            json=model_to_json_serializable_dict(get_artifact_response),
        )

        artifact = self.debusine.artifact_get(artifact_id)

        self.assert_token_key_included_in_all_requests(self.token)

        self.assertEqual(
            artifact,
            create_artifact_response(
                id=artifact_id,
                category=category,
                data=data,
                created_at=created_at,
                workspace=workspace,
                download_tar_gz_url=download_artifact_tar_gz_url,
            ),
        )

    @responses.activate
    def test_artifact_create_success(self):
        """Test artifact_create succeeds."""
        files = self.files_for_artifact()

        workspace = "test"
        work_request = 5
        expire_at = datetime.datetime.now() + datetime.timedelta(days=1)

        artifact_response = create_artifact_response(
            id=2,
            workspace=workspace,
            category="Testing",
            created_at=datetime.datetime.utcnow(),
            data={},
            files_to_upload=list(files.keys()),
        )

        responses.add(
            responses.POST,
            f"{self.api_url}/1.0/artifact/",
            json=model_to_json_serializable_dict(artifact_response),
            status=requests.codes.created,
        )

        artifact_to_post = WorkRequestDebugLogs(
            category=ArtifactCategory.WORK_REQUEST_DEBUG_LOGS,
            files=files,
            data={},
        )

        actual_artifact_response = self.debusine.artifact_create(
            artifact_to_post,
            workspace=workspace,
            work_request=work_request,
            expire_at=expire_at,
        )

        self.assertEqual(actual_artifact_response, artifact_response)

        self.assert_token_key_included_in_all_requests(self.token)

        self.assertEqual(
            json.loads(responses.calls[0].request.body),
            artifact_to_post.serialize_for_create_artifact(
                workspace=workspace,
                work_request=work_request,
                expire_at=expire_at,
            ),
        )

    @responses.activate
    def test_upload_files(self):
        """Client upload files to the server (Debusine.upload_files)."""
        file_relative = self.create_temporary_file(prefix="main.c")

        file_paths = [
            self.create_temporary_file(prefix="README"),
            self.create_temporary_file(prefix="AUTHORS?"),
            file_relative,
        ]

        artifact_id = 2

        urls = []
        for file_path in file_paths:
            file_path_encoded = urllib.parse.quote(file_path.name)
            url = (
                f"{self.api_url}/{Debusine.API_VERSION}/artifact/"
                f"{artifact_id}/files/{file_path_encoded}/"
            )
            responses.add(responses.PUT, url)
            urls.append(url)

        files_to_upload = {}
        for file_path in file_paths:
            files_to_upload[file_path.name] = file_path

        files_to_upload[file_relative.name] = file_relative.relative_to(
            file_relative.parent
        )

        self.debusine.upload_files(
            artifact_id, files_to_upload, file_relative.parent
        )

        for url in urls:
            self.assertTrue(responses.assert_call_count(url, 1))

    def test_upload_files_invalid_base_directory(self):
        """Base directory is None but there is a relative file path."""
        file_path = "src/main.c"

        files_to_upload = {"some/file/main.c": Path(file_path)}

        artifact_id = 2

        with self.assertRaisesRegex(
            ValueError,
            f"{file_path} is relative: base_directory parameter cannot be None",
        ):
            self.debusine.upload_files(artifact_id, files_to_upload, None)

    @responses.activate
    def test_upload_artifact(self):
        """upload_artifact() uploads the artifact and associated files."""
        category = ArtifactCategory.WORK_REQUEST_DEBUG_LOGS
        workspace = "the-workspace"
        work_request = 6

        data = {}
        filename = "README.txt"
        file_path = self.create_temporary_file()
        files = {
            filename: file_path,
            "already-uploaded-file.txt": self.create_temporary_file(),
        }

        local_artifact = WorkRequestDebugLogs(
            category=category, data=data, files=files
        )

        remote_id = 2

        # Create the artifact request
        responses.add(
            "POST",
            f"{self.debusine.api_url}/artifact/",
            create_artifact_response(
                id=remote_id,
                workspace=workspace,
                files_to_upload=[filename],
            ).json(),
            status=requests.codes.created,
        )

        # Upload the file request
        responses.add(
            "PUT",
            f"{self.debusine.api_url}/artifact/{remote_id}/files/{filename}/",
        )
        remote_artifact = self.debusine.upload_artifact(
            local_artifact, workspace=workspace, work_request=work_request
        )
        self.assertEqual(
            remote_artifact,
            RemoteArtifact(id=remote_id, workspace=workspace),
        )

        self.assertEqual(len(responses.calls), 2)

    @responses.activate
    def test_relation_create(self):
        """relation_create() makes an HTTP post."""
        artifact_id = 1
        target_id = 2
        relation_type = "extends"

        url = self.api_url + "/1.0/artifact-relation/"

        response_body = RelationResponse(
            id=7, artifact=artifact_id, target=target_id, type=relation_type
        )
        responses.add(
            responses.POST,
            url,
            json=dict(response_body),
        )

        self.debusine.relation_create(artifact_id, target_id, relation_type)

        responses.assert_call_count(url, 1)

        request_body = json.loads(responses.calls[0].request.body)

        self.assertEqual(
            request_body,
            RelationCreateRequest(
                artifact=artifact_id, target=target_id, type=relation_type
            ),
        )

    def add_get_download_artifact_response(self) -> ArtifactResponse:
        """
        Call responses.add("GET", "/api/1.0/artifact) with an Artifact.

        Return the ArtifactResponse that has been created.
        """
        artifact_id = 17
        url_download_artifact = (
            f"https://debusine.example.com/"
            f"artifact/{artifact_id}/?archive=tar.gz"
        )

        # Set endpoint to download artifact's information
        url_artifact_information = self.api_url + f"/1.0/artifact/{artifact_id}"
        artifact_response = create_artifact_response(
            id=artifact_id,
            download_tar_gz_url=url_download_artifact,
        )
        responses.add(
            "GET", url_artifact_information, body=artifact_response.json()
        )

        return artifact_response

    def download_artifact(
        self,
        server_body: io.BytesIO,
        artifact_id: int,
        destination: Path,
        *,
        tarball: bool,
    ):
        """Mock _get_fileobj_streaming, call self.debusine.download_artifact."""
        patcher = mock.patch.object(
            self.debusine, "_get_fileobj_streaming", autospec=True
        )
        mocked = patcher.start()
        mocked.return_value.raw = server_body
        self.addCleanup(patcher.stop)

        return self.debusine.download_artifact(
            artifact_id, destination, tarball=tarball
        )

    @responses.activate
    def test_download_artifact_failed(self):
        """download_artifact() download an invalid tar file Raise exception."""
        artifact_response = self.add_get_download_artifact_response()

        with self.assertRaisesRegex(
            RuntimeError, r"Error untarring https://.*\. Returncode: "
        ):
            self.download_artifact(
                io.BytesIO(b"Invalid tar"),
                artifact_response.id,
                self.create_temporary_directory(),
                tarball=False,
            )

    def create_tar(self) -> tuple[io.BytesIO, Path]:
        """
        Create a tar file with a file inside.

        :return: tuple [tarFile, Path of a file]
        """
        compressed = io.BytesIO()
        tar_file = tarfile.open(fileobj=compressed, mode="w|gz")

        file = self.create_temporary_file()
        tar_file.add(file, arcname=file.name)
        tar_file.close()

        compressed.seek(0)

        return compressed, file

    @responses.activate
    def test_download_artifact_untar(self):
        """download_artifact() download the artifact and untar it."""
        compressed, file = self.create_tar()
        downloaded_directory = self.create_temporary_directory()
        expected_artifact_response = self.add_get_download_artifact_response()

        stdout = io.StringIO()

        with contextlib.redirect_stdout(stdout):
            actual_artifact_response = self.download_artifact(
                compressed,
                expected_artifact_response.id,
                downloaded_directory,
                tarball=False,
            )

        self.assertEqual(
            file.read_bytes(),
            (downloaded_directory / file.name).read_bytes(),
        )

        self.assertEqual(actual_artifact_response, expected_artifact_response)

        self.logger.info.assert_has_calls(
            [
                call(
                    "Downloading artifact and uncompressing into %s",
                    downloaded_directory,
                ),
                call(file.name),
            ]
        )

    @responses.activate
    def test_download_artifact_tarball(self):
        """download_artifact() download the artifact as .tar.gz."""
        compressed, file = self.create_tar()
        downloaded_directory = self.create_temporary_directory()

        expected_artifact_response = self.add_get_download_artifact_response()
        artifact_id = expected_artifact_response.id

        self.download_artifact(
            compressed,
            expected_artifact_response.id,
            downloaded_directory,
            tarball=True,
        )

        artifact_file = downloaded_directory / f"artifact-{artifact_id}.tar.gz"

        self.assertEqual(compressed.getvalue(), artifact_file.read_bytes())

        self.logger.info.assert_called_with(
            "Artifact downloaded: %s", artifact_file
        )

    def add_get_download_artifact_file_response(
        self, contents
    ) -> ArtifactResponse:
        """
        Call responses.add("GET", "/api/1.0/artifact) with an Artifact.

        Return the ArtifactResponse that has been created.
        """
        artifact_id = 17
        url_system_tar_xz = (
            f"https://debusine.example.com/"
            f"artifact/{artifact_id}/system.tar.xz"
        )

        # Set endpoint to download artifact's information
        url_artifact_information = self.api_url + f"/1.0/artifact/{artifact_id}"
        artifact_response = create_artifact_response(
            id=artifact_id,
            files=FilesResponseType(
                {
                    "system.tar.xz": FileResponse(
                        size=len(contents),
                        checksums={
                            "sha256": hashlib.sha256(contents).hexdigest()
                        },
                        type="file",
                        url=pydantic.parse_obj_as(
                            pydantic.AnyUrl, url_system_tar_xz
                        ),
                    )
                }
            ),
        )
        responses.add(
            "GET", url_artifact_information, body=artifact_response.json()
        )

        return artifact_response

    def download_artifact_file(
        self,
        server_body: io.BytesIO,
        artifact_id: int,
        path_in_artifact: str,
        destination: Path,
    ):
        """Mock _get_fileobj_streaming, call debusine.download_artifact_file."""
        patcher = mock.patch.object(
            self.debusine, "_get_fileobj_streaming", autospec=True
        )
        mocked = patcher.start()
        mocked.return_value.iter_content.side_effect = [server_body]
        self.addCleanup(patcher.stop)

        return self.debusine.download_artifact_file(
            artifact_id, path_in_artifact, destination
        )

    @responses.activate
    def test_download_artifact_file_missing_path(self):
        """download_artifact_file() requires path_in_artifact to exist."""
        artifact_response = self.add_get_download_artifact_file_response(b"")

        with self.assertRaisesRegex(
            ValueError,
            f"No file 'nonexistent' in artifact {artifact_response.id}",
        ):
            self.download_artifact_file(
                io.BytesIO(),
                artifact_response.id,
                "nonexistent",
                self.create_temporary_directory() / "nonexistent",
            )

    @responses.activate
    def test_download_artifact_file(self):
        """download_artifact_file() downloads the requested file."""
        file = io.BytesIO(b"File contents")
        downloaded_directory = self.create_temporary_directory()
        artifact_response = self.add_get_download_artifact_file_response(
            file.getvalue()
        )
        destination = downloaded_directory / "system.tar.xz"

        self.download_artifact_file(
            file, artifact_response.id, "system.tar.xz", destination
        )

        self.assertEqual(destination.read_bytes(), file.getvalue())
        self.logger.info.assert_called_with(
            "Artifact file downloaded: %s", destination
        )

    @responses.activate
    def test_lookup_single(self):
        """lookup_single() performs the requested lookup."""
        responses.add(
            responses.POST,
            f"{self.debusine.api_url}/lookup/single/",
            json={
                "result_type": CollectionItemType.ARTIFACT,
                "collection_item": None,
                "artifact": 1,
                "collection": None,
            },
        )

        response = self.debusine.lookup_single(
            lookup="collection/artifact",
            work_request=1,
            expect_type=CollectionItemType.ARTIFACT,
            default_category="test",
        )

        self.assertEqual(
            response,
            LookupSingleResponse(
                result_type=CollectionItemType.ARTIFACT, artifact=1
            ),
        )
        self.assertEqual(
            json.loads(responses.calls[0].request.body),
            {
                "lookup": "collection/artifact",
                "work_request": 1,
                "default_category": "test",
                "expect_type": CollectionItemType.ARTIFACT,
            },
        )

    @responses.activate
    def test_lookup_multiple(self):
        """lookup_multiple() performs the requested lookup."""
        responses.add(
            responses.POST,
            f"{self.debusine.api_url}/lookup/multiple/",
            json=[
                {
                    "result_type": CollectionItemType.ARTIFACT,
                    "collection_item": artifact_id,
                    "artifact": artifact_id,
                    "collection": None,
                }
                for artifact_id in (1, 2)
            ],
        )

        response = self.debusine.lookup_multiple(
            lookup=LookupMultiple.parse_obj({"collection": "collection"}),
            work_request=1,
            expect_type=CollectionItemType.ARTIFACT,
            default_category="test",
        )

        self.assertEqual(
            response,
            LookupMultipleResponse.parse_obj(
                [
                    LookupSingleResponse(
                        result_type=CollectionItemType.ARTIFACT,
                        collection_item=artifact_id,
                        artifact=artifact_id,
                    )
                    for artifact_id in (1, 2)
                ]
            ),
        )
        self.assertEqual(
            json.loads(responses.calls[0].request.body),
            {
                # The lookup is normalized.
                "lookup": [
                    {
                        "category": None,
                        "child_type": "artifact",
                        "collection": "collection",
                        "data_matchers": [],
                        "name_matcher": None,
                    }
                ],
                "work_request": 1,
                "default_category": "test",
                "expect_type": CollectionItemType.ARTIFACT,
            },
        )

    def test_get_fileobj_streaming(self):
        """Test requests.get() is called as expected."""
        patcher = mock.patch("requests.sessions.Session.get", autospec=True)
        mocker = patcher.start()

        expected_raw = b"test"
        mocker.return_value = mock.create_autospec(
            spec=requests.models.Response, raw=expected_raw
        )
        self.addCleanup(patcher.stop)

        url = "https://debian.org"
        headers = {"User-Agent": "debusine"}

        r = self.debusine._get_fileobj_streaming(url, headers=headers)

        self.assertEqual(r.raw, expected_raw)

        mocker.assert_called_with(mock.ANY, url, headers=headers, stream=True)


class DebusineAsyncTests(TestHelpersMixin, server.DebusineAioHTTPTestCase):
    """
    Tests for Debusine client asynchronous functionality.

    Use server.DebusineAioHTTPTestCase to mock the server (instead of the
    responses library used by DebusineTests).
    """

    async def setUpAsync(self):
        """Set up tests."""
        await super().setUpAsync()

        self.api_url = str(self.server.make_url("/api"))

        self.token = "token"
        self.logger = mock.create_autospec(spec=logging.Logger)
        self.debusine = Debusine(self.api_url, self.token, logger=self.logger)

    def on_complete_url(self, *, params: dict[str, str] | None = None) -> str:
        """Return the WebSocket on-completed-url."""
        url = self.api_url + "/ws/1.0/work-request/on-completed/"
        url = url.replace("http", "ws", 1)

        if params is not None:
            url += "?" + urlencode(params)

        return url

    def script_write_parameters_to_file(self) -> tuple[Path, Path]:
        """
        Create a script that save its parameters into a file.

        :return: tuple with Path to the script and Path to the output file.

        """
        output_file = self.create_temporary_file()

        command = self.create_temporary_file(
            contents=textwrap.dedent(
                f"""\
            #!/bin/bash
            echo $* >> {str(output_file)}
            """
            ).encode("utf-8")
        )

        os.chmod(command, stat.S_IXUSR | stat.S_IRUSR)

        return command, output_file

    async def test_wait_and_execute_command(self):
        """A WorkRequest finished and command is executed."""
        await self.execute_wait_and_execute(
            command_modifier=lambda cmd: cmd,
            working_directory=lambda _: Path.cwd(),  # noqa: U101
            on_completed_work_requests=1,
        )
        self.assertEqual(self.requests[-1].headers["Token"], self.token)

    async def test_wait_and_execute_two_commands(self):
        """Two WorkRequests finished: two commands are executed."""
        await self.execute_wait_and_execute(
            command_modifier=lambda cmd: cmd,
            working_directory=lambda _: Path.cwd(),  # noqa: U101
            on_completed_work_requests=2,
        )

    async def test_wait_and_execute_command_relative_to_working_directory(self):
        """A WorkRequest finished and command is executed."""
        await self.execute_wait_and_execute(
            command_modifier=lambda cmd: f"./{cmd.name}",
            working_directory=lambda cmd: cmd.parent,
            on_completed_work_requests=1,
        )

    async def execute_wait_and_execute(
        self,
        command_modifier: Callable[[Path], Path],
        working_directory: Callable[[Path], Path],
        on_completed_work_requests: int,
    ):
        """
        Test self.debusine._wait_and_execute().

        Create a script that use as a command. Execute and assert that
        it was executed.

        :param command_modifier: how to modify the Path to the script before
          executing it (e.g. to convert it to a relative path)
        :param working_directory: which working directory to use (e.g.
          to use the script's working directory or Path.cwd().
        :param on_completed_work_requests: number of on completed work requests
          expected. The IDs will increment by one on each finished work request.
        """
        url = self.on_complete_url()
        command, output_file = self.script_write_parameters_to_file()

        modified_command = command_modifier(command)

        self.ON_COMPLETED_WORK_REQUESTS = on_completed_work_requests

        await self.debusine._wait_and_execute(
            url=url,
            last_completed_at=None,
            command=modified_command,
            working_directory=working_directory(command),
        )

        expected_in_file = ""
        expected_log_executing_calls = []

        for number in range(0, on_completed_work_requests):
            expected_in_file += (
                f"{server.DebusineAioHTTPTestCase.WORK_REQUEST_ID + number} "
                f"{server.DebusineAioHTTPTestCase.WORK_REQUEST_RESULT}\n"
            )

            cmd = [
                str(modified_command),
                str(server.DebusineAioHTTPTestCase.WORK_REQUEST_ID + number),
                str(server.DebusineAioHTTPTestCase.WORK_REQUEST_RESULT),
            ]

            expected_log_executing_calls.append(call("Executing %s", cmd))

        self.assertEqual(output_file.read_text(), expected_in_file)

        self.logger.info.assert_has_calls(
            [
                call("Requesting %s", url),
                call("Connected!"),
                *expected_log_executing_calls,
            ],
        )

    async def test_wait_and_execute_command_write_last_completed_file(self):
        """wait_and_execute_command update last_completed_at file."""
        last_completed_at = self.create_temporary_file()

        command, _ = self.script_write_parameters_to_file()

        url = self.on_complete_url()

        await self.debusine._wait_and_execute(
            url=url,
            last_completed_at=last_completed_at,
            command=str(command),
            working_directory=Path.cwd(),
        )

        completed_at = DebusineAioHTTPTestCase.COMPLETED_AT.isoformat()
        expected = (
            json.dumps({"last_completed_at": completed_at}, indent=2) + "\n"
        )

        self.assertEqual(last_completed_at.read_text(), expected)

    def patch_tenacity_sleep(self, attempts: int = 1) -> MagicMock:
        """Patch self.debusine._tenacity_sleep and return its mock."""
        patcher = mock.patch.object(
            self.debusine, "_tenacity_sleep", autospec=True
        )
        mocked = patcher.start()
        mocked.side_effect = [None] * attempts
        self.addCleanup(patcher.stop)

        return mocked

    async def test_wait_and_execute_command_retry(self):
        """Client get disconnected: try again."""
        attempts = 2
        tenacity_sleep_mocked = self.patch_tenacity_sleep(attempts)

        exc_message = "SSL error"

        exc = aiohttp.client_exceptions.ClientError(exc_message)

        # Mock ws_connect to raise the ClientError exception
        def mock_ws_connect(*args, **kwargs):
            class MockWebSocketResponse:
                async def __aenter__(self):
                    raise exc

                async def __aexit__(self, exc_type, exc_val, exc_tb):
                    pass  # pragma: no cover

            return MockWebSocketResponse()

        patcher = mock.patch.object(
            aiohttp.ClientSession, 'ws_connect', new=mock_ws_connect
        )
        patcher.start()
        self.addCleanup(patcher.stop)

        with self.assertRaises(StopAsyncIteration):
            await self.debusine._wait_and_execute(
                url=self.on_complete_url(),
                command="/bin/echo",
                working_directory=Path.cwd(),
                last_completed_at=None,
            )

        # Called three times: twice returned None and once StopAsyncIteration
        self.assertEqual(tenacity_sleep_mocked.call_count, attempts + 1)

        self.logger.error.assert_has_calls([call("  Error: %s", exc)] * 3)

    async def test_wait_and_execute_command_fail(self):
        """Command cannot be executed: raise FileNotFoundError."""
        with self.assertRaises(FileNotFoundError):
            await self.debusine._wait_and_execute(
                url=self.on_complete_url(),
                command="file-does-not-exist",
                working_directory=Path.cwd(),
                last_completed_at=None,
            )

    async def test_wait_and_execute_no_token(self):
        """_wait_and_execute works even if it has no token."""
        self.debusine = Debusine(self.api_url, None, logger=self.logger)
        await self.execute_wait_and_execute(
            command_modifier=lambda cmd: cmd,
            working_directory=lambda _: Path.cwd(),  # noqa: U101
            on_completed_work_requests=1,
        )
        self.assertNotIn("Token", self.requests[-1].headers)

    async def test_tenacity_sleep(self):
        """Debusine._tenacity_sleep call asyncio.sleep."""
        patched = mock.patch("asyncio.sleep", autospec=True)
        mocked = patched.start()
        self.addCleanup(patched.stop)

        delay = 1.5
        await self.debusine._tenacity_sleep(delay)

        mocked.assert_called_with(delay)

    def patch_wait_and_execute(self) -> MagicMock:
        """Patch self.debusine._wait_and_execute, return its mock."""
        patcher = mock.patch.object(
            self.debusine, "_wait_and_execute", return_value=None
        )
        mocked = patcher.start()
        self.addCleanup(patcher.stop)

        return mocked

    def test_on_work_request_completed(self):
        """on_work_request_completed call _wait_and_executed. No workspaces."""
        command = "this-is-the-command"

        wait_and_executed_mocked = self.patch_wait_and_execute()

        self.debusine.on_work_request_completed(
            workspaces=None,
            command=command,
            working_directory=Path.cwd(),
            last_completed_at=None,
        )

        url = self.on_complete_url()

        wait_and_executed_mocked.assert_awaited_with(
            url=url,
            command=command,
            working_directory=Path.cwd(),
            last_completed_at=None,
        )

    def test_on_work_request_completed_with_workspaces(self):
        """on_work_request_completed call _wait_and_execute with workspaces."""
        workspace = "lts"
        command = "this-is-the-command"

        wait_and_executed_mocked = self.patch_wait_and_execute()

        self.debusine.on_work_request_completed(
            workspaces=[workspace],
            last_completed_at=None,
            command=command,
            working_directory=Path.cwd(),
        )

        url = self.on_complete_url(params={"workspaces": workspace})

        wait_and_executed_mocked.assert_awaited_with(
            url=url,
            command=command,
            working_directory=Path.cwd(),
            last_completed_at=None,
        )

    def test_on_work_request_completed_last_completed_at(self):
        """
        on_work_request_completed call _wait_and_execute with last_completed.

        The file in last_completed_at existed and its contents is read.
        """
        command = "this-is-the-command"
        last_completed_at = datetime.datetime.utcnow().isoformat()

        last_completed_at_file = self.create_temporary_file(
            contents=json.dumps(
                {"last_completed_at": last_completed_at}
            ).encode("utf-8")
        )

        wait_and_executed_mocked = self.patch_wait_and_execute()

        self.debusine.on_work_request_completed(
            last_completed_at=last_completed_at_file,
            command=command,
            working_directory=Path.cwd(),
        )

        url = self.on_complete_url(
            params={"completed_at_since": last_completed_at}
        )

        wait_and_executed_mocked.assert_called_with(
            url=url,
            command=command,
            working_directory=Path.cwd(),
            last_completed_at=last_completed_at_file,
        )

    def test_on_work_request_completed_last_completed_at_does_not_exist(self):
        """
        on_work_request_completed call _wait_and_execute with last_completed.

        The file in last_completed_at did not exist.
        """
        last_completed_at_file = self.create_temporary_file()
        last_completed_at_file.unlink()

        command = "this-is-the-command"

        wait_and_executed_mocked = self.patch_wait_and_execute()

        self.debusine.on_work_request_completed(
            last_completed_at=last_completed_at_file,
            command=command,
            working_directory=Path.cwd(),
        )

        url = self.on_complete_url()

        wait_and_executed_mocked.assert_called_with(
            url=url,
            command=command,
            working_directory=Path.cwd(),
            last_completed_at=last_completed_at_file,
        )
