# Copyright © 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.

"""Tests for base workflow code."""

import logging
from unittest import mock

from debusine.artifacts.models import ArtifactCategory
from debusine.db.models import ArtifactRelation, WorkRequest
from debusine.server.workflows.base import (
    Workflow,
    WorkflowRunError,
    orchestrate_workflow,
)
from debusine.server.workflows.models import BaseWorkflowData
from debusine.server.workflows.noop import NoopWorkflow
from debusine.tasks import TaskConfigError
from debusine.tasks.models import BaseDynamicTaskData, TaskTypes
from debusine.tasks.server import TaskDatabaseInterface
from debusine.tasks.tests.helper_mixin import TestTaskMixin
from debusine.test.django import TestCase
from debusine.test.utils import preserve_task_registry


class WorkflowClassTests(TestCase):
    """Test the Workflow base class."""

    def test_ensure_artifact_categories_matches(self) -> None:
        """Test with a matching artifact."""
        artifact, _ = self.playground.create_artifact(
            category=ArtifactCategory.TEST
        )
        Workflow.ensure_artifact_categories(
            artifact, [ArtifactCategory.TEST, ArtifactCategory.UPLOAD]
        )

    def test_ensure_artifact_categories_mismatches(self) -> None:
        """Test with a mis-matching artifact."""
        artifact, _ = self.playground.create_artifact(
            category=ArtifactCategory.TEST
        )
        with self.assertRaisesRegex(
            TaskConfigError,
            r'Category of ".* debusine:test .*" is not one of '
            r"\['debian:system-tarball', 'debian:system-image'\]",
        ):
            Workflow.ensure_artifact_categories(
                artifact,
                [
                    ArtifactCategory.SYSTEM_TARBALL,
                    ArtifactCategory.SYSTEM_IMAGE,
                ],
            )

    def test_follow_artifact_relation_upload(self) -> None:
        """Test finding a source package from an upload."""
        upload_artifacts = self.playground.create_upload_artifacts()
        test_artifact, _ = self.playground.create_artifact()

        # Create some extra unrelated relations
        self.playground.create_artifact_relation(
            artifact=upload_artifacts.upload,
            target=upload_artifacts.binaries[0],
            relation_type=ArtifactRelation.Relations.BUILT_USING,
        )
        self.playground.create_artifact_relation(
            artifact=upload_artifacts.upload,
            target=test_artifact,
            relation_type=ArtifactRelation.Relations.EXTENDS,
        )

        artifact = Workflow.follow_artifact_relation(
            upload_artifacts.upload,
            ArtifactRelation.Relations.EXTENDS,
            ArtifactCategory.SOURCE_PACKAGE,
        )
        self.assertEqual(artifact, upload_artifacts.source)

    def test_follow_artifact_relation_not_found(self) -> None:
        """Test with a lone debian:upload (no relations)."""
        upload_artifacts = self.playground.create_upload_artifacts(source=False)

        with self.assertRaisesRegex(
            TaskConfigError,
            r"Unable to find an artifact of category debian:source-package "
            r'with a relationship of type extends from ".* debian:upload .*"',
        ):
            Workflow.follow_artifact_relation(
                upload_artifacts.upload,
                ArtifactRelation.Relations.EXTENDS,
                ArtifactCategory.SOURCE_PACKAGE,
            )

    def test_follow_artifact_relation_multiple_found(self) -> None:
        """Test with a debian:upload extending multiple srcpkgs."""
        upload_artifacts = self.playground.create_upload_artifacts(source=True)
        source2 = self.playground.create_source_artifact()
        self.playground.create_artifact_relation(
            artifact=upload_artifacts.upload,
            target=source2,
            relation_type=ArtifactRelation.Relations.EXTENDS,
        )

        with self.assertRaisesRegex(
            TaskConfigError,
            r"Multiple artifacts of category debian:source-package with "
            r'a relationship of type extends from ".* debian:upload .*" found',
        ):
            Workflow.follow_artifact_relation(
                upload_artifacts.upload,
                ArtifactRelation.Relations.EXTENDS,
                ArtifactCategory.SOURCE_PACKAGE,
            )

    def test_locate_debian_source_package_with_srcpkg(self) -> None:
        """Test with a debian:source-package."""
        source_artifact = self.playground.create_source_artifact()
        result = Workflow.locate_debian_source_package(source_artifact)
        self.assertEqual(result, source_artifact)

    def test_locate_debian_source_package_with_upload(self) -> None:
        """Test with a debian:source-package."""
        upload_artifacts = self.playground.create_upload_artifacts()
        result = Workflow.locate_debian_source_package(upload_artifacts.upload)
        self.assertEqual(result, upload_artifacts.source)

    def test_locate_debian_source_package_with_junk(self) -> None:
        """Test with a debian:source-package."""
        artifact, _ = self.playground.create_artifact(
            category=ArtifactCategory.TEST
        )
        with self.assertRaisesRegex(
            TaskConfigError,
            r'Category of ".* debusine:test .*" is not one of '
            r"\['debian:source-package', 'debian:upload'\]",
        ):
            Workflow.locate_debian_source_package(artifact)


class WorkflowInstanceTests(TestCase):
    """Test instances of the Workflow class."""

    def setUp(self) -> None:
        """Create a NoopWorkflow for testing."""
        super().setUp()
        w = self.playground.create_workflow(task_name="noop")
        self.workflow = NoopWorkflow(w)

    def test_locate_debian_source_package_lookup_srcpkg(self) -> None:
        """Test locate_debian_source_package_lookup() with a source package."""
        source_artifact = self.playground.create_source_artifact()
        self.assertEqual(
            self.workflow.locate_debian_source_package_lookup(
                source_artifact.id
            ),
            source_artifact.id,
        )

    def test_locate_debian_source_package_lookup_upload(self) -> None:
        """Test locate_debian_source_package_lookup() with a debian:upload."""
        artifacts = self.playground.create_upload_artifacts()
        self.assertEqual(
            self.workflow.locate_debian_source_package_lookup(
                artifacts.upload.id
            ),
            f"{artifacts.source.id}@artifacts",
        )


class OrchestrateWorkflowTests(TestCase):
    """Test orchestrate_workflow()."""

    def test_workflow_callback(self) -> None:
        """A workflow callback is run and marked as completed."""
        parent = self.playground.create_workflow(task_name="noop")
        parent.mark_running()
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test"
        )

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.callback"
        ) as mock_noop_callback:
            orchestrate_workflow(wr)

        mock_noop_callback.assert_called_once_with(wr)
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.SUCCESS)

    @preserve_task_registry()
    def test_workflow_callback_computes_dynamic_data(self) -> None:
        """Dynamic task data is computed for workflow callbacks."""

        class ExampleDynamicData(BaseDynamicTaskData):
            dynamic: str

        class ExampleWorkflow(
            TestTaskMixin, Workflow[BaseWorkflowData, ExampleDynamicData]
        ):
            TASK_NAME = "example"

            def compute_dynamic_data(
                self, task_database: TaskDatabaseInterface  # noqa: U100
            ) -> ExampleDynamicData:
                return ExampleDynamicData(dynamic="foo")

            def populate(self) -> None:
                """Unused abstract method from Workflow."""
                raise NotImplementedError()

        parent = self.playground.create_workflow(task_name="example")
        parent.mark_running()
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test"
        )

        with mock.patch.object(ExampleWorkflow, "callback") as mock_callback:
            orchestrate_workflow(wr)

        mock_callback.assert_called_once_with(wr)
        parent.refresh_from_db()
        self.assertEqual(parent.dynamic_task_data, {"dynamic": "foo"})
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.SUCCESS)

    @preserve_task_registry()
    def test_workflow_callback_computes_dynamic_data_only_once(self) -> None:
        """If a workflow already has dynamic data, it is left alone."""

        class ExampleDynamicData(BaseDynamicTaskData):
            dynamic: str

        class ExampleWorkflow(
            TestTaskMixin, Workflow[BaseWorkflowData, ExampleDynamicData]
        ):
            TASK_NAME = "example"

            def populate(self) -> None:
                """Unused abstract method from Workflow."""
                raise NotImplementedError()

        WorkRequest.objects.all().delete()
        parent = self.playground.create_workflow(task_name="example")
        parent.dynamic_task_data = {"dynamic": "foo"}
        parent.save()
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test"
        )

        with (
            mock.patch.object(
                ExampleWorkflow, "compute_dynamic_data"
            ) as mock_compute_dynamic_data,
            mock.patch.object(ExampleWorkflow, "callback") as mock_callback,
        ):
            orchestrate_workflow(wr)

        mock_compute_dynamic_data.assert_not_called()
        mock_callback.assert_called_once_with(wr)
        parent.refresh_from_db()
        self.assertEqual(parent.dynamic_task_data, {"dynamic": "foo"})
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.SUCCESS)

    def test_workflow_callback_outside_workflow(self) -> None:
        """A workflow callback outside a workflow is skipped."""
        wr = self.playground.create_work_request(
            task_type=TaskTypes.INTERNAL, task_name="workflow"
        )
        expected_message = "Workflow callback is not contained in a workflow"

        with (
            self.assertLogsContains(
                f"Error running work request Internal/workflow ({wr.id}): "
                f"{expected_message}",
                logger="debusine.server.workflows.base",
                level=logging.WARNING,
            ),
            self.assertRaises(WorkflowRunError) as raised,
        ):
            orchestrate_workflow(wr)

        self.assertEqual(raised.exception.args, (wr, expected_message))
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)

    def test_workflow_callback_bad_task_data(self) -> None:
        """A workflow callback with bad task data is skipped."""
        parent = self.playground.create_workflow(
            task_name="noop", task_data={"nonsense": ""}, validate=False
        )
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test"
        )
        expected_message = "failed to configure: 1 validation error"

        with (
            mock.patch(
                "debusine.server.workflows.noop.NoopWorkflow.callback"
            ) as mock_noop_callback,
            self.assertLogsContains(
                f"Error running work request Internal/workflow ({wr.id}): "
                f"{expected_message}",
                logger="debusine.server.workflows.base",
                level=logging.WARNING,
            ),
            self.assertRaises(WorkflowRunError) as raised,
        ):
            orchestrate_workflow(wr)

        mock_noop_callback.assert_not_called()
        self.assertEqual(len(raised.exception.args), 2)
        self.assertEqual(raised.exception.args[0], wr)
        self.assertTrue(raised.exception.args[1].startswith(expected_message))
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)

    def test_workflow_callback_with_failing_compute_dynamic_data(self) -> None:
        """A workflow callback where computing dynamic data fails is skipped."""
        parent = self.playground.create_workflow(task_name="noop")
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test"
        )

        with (
            mock.patch(
                "debusine.server.workflows.noop.NoopWorkflow"
                ".compute_dynamic_data",
                side_effect=Exception("Boom"),
            ),
            mock.patch(
                "debusine.server.workflows.noop.NoopWorkflow.callback"
            ) as mock_noop_callback,
            self.assertLogs("debusine", level=logging.WARNING) as log,
            self.assertRaises(WorkflowRunError) as raised,
        ):
            orchestrate_workflow(wr)

        mock_noop_callback.assert_not_called()
        self.assertIn(
            f"Workflow orchestrator failed to pre-process task data for "
            f"WorkRequest {parent.id}. "
            f"Task data: {parent.task_data} Error: Boom",
            "\n".join(log.output),
        )
        expected_message = "failed to compute dynamic data"
        self.assertIn(
            f"Error running work request Internal/workflow ({wr.id}): "
            f"{expected_message}",
            "\n".join(log.output),
        )
        self.assertEqual(raised.exception.args, (wr, expected_message))
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)

    def test_workflow_callback_fails(self) -> None:
        """A workflow callback that fails is logged."""
        parent = self.playground.create_workflow(task_name="noop")
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test"
        )
        expected_message = "orchestrator failed: Boom"

        with (
            mock.patch(
                "debusine.server.workflows.noop.NoopWorkflow.callback",
                side_effect=ValueError("Boom"),
            ) as mock_noop_callback,
            self.assertLogsContains(
                f"Error running work request Internal/workflow ({wr.id}): "
                f"{expected_message}",
                logger="debusine.server.workflows.base",
                level=logging.WARNING,
            ),
            self.assertRaises(WorkflowRunError) as raised,
        ):
            orchestrate_workflow(wr)

        mock_noop_callback.assert_called_once_with(wr)
        self.assertEqual(raised.exception.args, (wr, expected_message))
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)

    def test_workflow(self) -> None:
        """A workflow is populated and left running."""
        wr = self.playground.create_workflow(task_name="noop")
        wr.mark_running()

        def populate() -> None:
            wr.create_child("noop")

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.populate",
            side_effect=populate,
        ) as mock_noop_populate:
            orchestrate_workflow(wr)

        mock_noop_populate.assert_called_once_with()
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.RUNNING)
        self.assertEqual(wr.result, WorkRequest.Results.NONE)

    def test_workflow_empty(self) -> None:
        """An empty workflow is populated and marked as completed."""
        wr = self.playground.create_workflow(task_name="noop")
        wr.mark_running()

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.populate"
        ) as mock_noop_populate:
            orchestrate_workflow(wr)

        mock_noop_populate.assert_called_once_with()
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.SUCCESS)

    @preserve_task_registry()
    def test_workflow_computes_dynamic_data(self) -> None:
        """Dynamic task data is computed for workflows."""

        class ExampleDynamicData(BaseDynamicTaskData):
            dynamic: str

        class ExampleWorkflow(
            TestTaskMixin, Workflow[BaseWorkflowData, ExampleDynamicData]
        ):
            TASK_NAME = "example"

            def compute_dynamic_data(
                self, task_database: TaskDatabaseInterface  # noqa: U100
            ) -> ExampleDynamicData:
                return ExampleDynamicData(dynamic="foo")

            def populate(self) -> None:
                self.work_request.create_child("noop")

        wr = self.playground.create_workflow(task_name="example")
        wr.mark_running()

        orchestrate_workflow(wr)

        wr.refresh_from_db()
        self.assertEqual(wr.dynamic_task_data, {"dynamic": "foo"})
        self.assertEqual(wr.status, WorkRequest.Statuses.RUNNING)
        self.assertEqual(wr.result, WorkRequest.Results.NONE)

    def test_workflow_bad_task_data(self) -> None:
        """A workflow with bad task data is skipped."""
        # Bad task data would normally be caught by create_workflow.  Force
        # it to happen here by changing the task data after initial
        # creation.  (A more realistic case might be one where the
        # definition of a workflow changes but existing data isn't
        # migrated.)
        wr = self.playground.create_workflow(task_name="noop")
        wr.task_data = {"nonsense": ""}
        wr.save()
        wr.mark_running()
        expected_message = "failed to configure: 1 validation error"

        with (
            mock.patch(
                "debusine.server.workflows.noop.NoopWorkflow.populate"
            ) as mock_noop_populate,
            self.assertLogsContains(
                f"Error running work request Workflow/noop ({wr.id}): "
                f"{expected_message}",
                logger="debusine.server.workflows.base",
                level=logging.WARNING,
            ),
            self.assertRaises(WorkflowRunError) as raised,
        ):
            orchestrate_workflow(wr)

        mock_noop_populate.assert_not_called()
        self.assertEqual(len(raised.exception.args), 2)
        self.assertEqual(raised.exception.args[0], wr)
        self.assertTrue(raised.exception.args[1].startswith(expected_message))
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)

    def test_workflows_with_failing_compute_dynamic_data(self) -> None:
        """A workflow where computing dynamic data fails is skipped."""
        wr = self.playground.create_workflow(task_name="noop")
        wr.mark_running()

        with (
            mock.patch(
                "debusine.server.workflows.noop.NoopWorkflow"
                ".compute_dynamic_data",
                side_effect=Exception("Boom"),
            ),
            mock.patch(
                "debusine.server.workflows.noop.NoopWorkflow.populate"
            ) as mock_noop_populate,
            self.assertLogs("debusine", level=logging.WARNING) as log,
            self.assertRaises(WorkflowRunError) as raised,
        ):
            orchestrate_workflow(wr)

        mock_noop_populate.assert_not_called()
        self.assertIn(
            f"Workflow orchestrator failed to pre-process task data for "
            f"WorkRequest {wr.id}. "
            f"Task data: {wr.task_data} Error: Boom",
            "\n".join(log.output),
        )
        expected_message = "failed to compute dynamic data"
        self.assertIn(
            f"Error running work request Workflow/noop ({wr.id}): "
            f"{expected_message}",
            "\n".join(log.output),
        )
        self.assertEqual(raised.exception.args, (wr, expected_message))
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)

    def test_workflow_fails(self) -> None:
        """A workflow that fails is logged."""
        wr = self.playground.create_workflow(task_name="noop")
        wr.mark_running()
        expected_message = "orchestrator failed: Boom"

        with (
            mock.patch(
                "debusine.server.workflows.noop.NoopWorkflow.populate",
                side_effect=ValueError("Boom"),
            ) as mock_noop_populate,
            self.assertLogsContains(
                f"Error running work request Workflow/noop ({wr.id}): "
                f"{expected_message}",
                logger="debusine.server.workflows.base",
                level=logging.WARNING,
            ),
            self.assertRaises(WorkflowRunError) as raised,
        ):
            orchestrate_workflow(wr)

        mock_noop_populate.assert_called_once_with()
        self.assertEqual(raised.exception.args, (wr, expected_message))
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)

    def test_workflow_unblocks_children(self) -> None:
        """Workflow children are unblocked if possible."""
        wr = self.playground.create_workflow(task_name="noop")
        wr.mark_running()

        def populate() -> None:
            children = [wr.create_child("noop") for _ in range(2)]
            children[1].add_dependency(children[0])

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.populate",
            side_effect=populate,
        ) as mock_noop_populate:
            orchestrate_workflow(wr)

        mock_noop_populate.assert_called_once_with()
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.RUNNING)
        self.assertEqual(wr.result, WorkRequest.Results.NONE)
        children = wr.children.order_by("id")
        self.assertEqual(children.count(), 2)
        self.assertEqual(children[0].status, WorkRequest.Statuses.PENDING)
        self.assertEqual(children[1].status, WorkRequest.Statuses.BLOCKED)

    def test_wrong_task_type(self) -> None:
        """Attempts to orchestrate non-workflows are logged."""
        wr = self.playground.create_work_request()
        expected_message = "does not have a workflow orchestrator"

        with (
            self.assertLogsContains(
                f"Error running work request {wr.task_type}/{wr.task_name} "
                f"({wr.id}): {expected_message}",
                logger="debusine.server.workflows.base",
                level=logging.WARNING,
            ),
            self.assertRaises(WorkflowRunError) as raised,
        ):
            orchestrate_workflow(wr)

        self.assertEqual(raised.exception.args, (wr, expected_message))
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)
