# 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 functools import partial
from typing import ClassVar, override
from unittest import mock

import debusine.tasks.tags as ttags
from debusine.artifacts.models import WorkRequestResults
from debusine.client.models import RuntimeParameter
from debusine.db.context import context
from debusine.db.models import WorkRequest
from debusine.db.models.lookups import LookupResultArtifact
from debusine.db.playground import scenarios
from debusine.server.tasks.wait.models import ExternalDebsignData
from debusine.server.workflows.base import (
    Workflow,
    WorkflowValidationError,
    orchestrate_workflow,
)
from debusine.server.workflows.models import (
    BaseWorkflowData,
    WorkRequestWorkflowData,
)
from debusine.server.workflows.tests.helpers import SampleWorkflow
from debusine.tasks.models import (
    BaseDynamicTaskData,
    OutputData,
    OutputDataError,
)
from debusine.test.django import TestCase, preserve_db_task_registry


class WorkflowTests(TestCase):
    """Test :py:class:`Workflow`."""

    def make_concrete_workflow(
        self,
    ) -> Workflow[BaseWorkflowData, BaseDynamicTaskData]:
        """Instantiate a concrete instance of Workflow."""

        class ConcreteWorkflow(Workflow[BaseWorkflowData, BaseDynamicTaskData]):
            """Concrete implementation of Workflow."""

            TASK_VERSION = 1

            @override
            def populate(self) -> None:
                raise NotImplementedError()

            @override
            def compute_dynamic_data(self) -> BaseDynamicTaskData:
                return BaseDynamicTaskData()

            @override
            def _execute(self) -> WorkRequestResults:
                raise NotImplementedError()

        work_request = self.playground.create_workflow(
            task_name="concreteworkflow"
        )
        task = work_request.get_task()
        assert isinstance(task, Workflow)
        return task

    @preserve_db_task_registry()
    def test_get_label_no_workflow_data_no_dynamic_data(self) -> None:
        task = self.make_concrete_workflow()
        task.work_request.workflow_data_json = {}
        self.assertEqual(task.get_label(), "concreteworkflow")

    @preserve_db_task_registry()
    def test_get_label_workflow_data_no_dynamic_data(self) -> None:
        task = self.make_concrete_workflow()
        task.work_request.workflow_data_json = {
            "workflow_template_name": "concreteworkflow-template"
        }
        self.assertEqual(task.get_label(), "concreteworkflow-template")

    @preserve_db_task_registry()
    def test_get_label_no_parameter_summary(self) -> None:
        task = self.make_concrete_workflow()
        task.ensure_dynamic_data()
        self.assertEqual(task.get_label(), "concreteworkflow")

    @preserve_db_task_registry()
    def test_get_label_parameter_summary(self) -> None:
        task = self.make_concrete_workflow()
        task.ensure_dynamic_data()
        assert task.dynamic_data is not None
        task.dynamic_data.parameter_summary = "summary"
        self.assertEqual(task.get_label(), "concreteworkflow: summary")

    def test_get_label_callback_no_wf_data_no_dyn_data(self) -> None:
        template = self.playground.create_workflow_template(
            name="test", task_name="noop"
        )
        workflow = self.playground.create_workflow(template, task_data={})
        workflow.workflow_data_json = {}
        workflow.save()
        wr = WorkRequest.objects.create_workflow_callback(
            parent=workflow, step="test", display_name="Test"
        )
        wr.workflow_data_json = {}
        wr.save()
        self.assertEqual(wr.get_task().get_label(), "noop")

    def test_get_label_callback_no_dyn_data_no_display_name(self) -> None:
        template = self.playground.create_workflow_template(
            name="test", task_name="noop"
        )
        workflow = self.playground.create_workflow(template, task_data={})
        wr = WorkRequest.objects.create_workflow_callback(
            parent=workflow, step="test-step"
        )
        self.assertEqual(wr.get_task().get_label(), "test-step")

    def test_get_label_callback_no_dyn_data(self) -> None:
        template = self.playground.create_workflow_template(
            name="test", task_name="noop"
        )
        workflow = self.playground.create_workflow(template, task_data={})
        wr = WorkRequest.objects.create_workflow_callback(
            parent=workflow, step="test-step", display_name="Test Step"
        )
        self.assertEqual(wr.get_task().get_label(), "Test Step")

    def test_get_label_callback_parameter_summary(self) -> None:
        template = self.playground.create_workflow_template(
            name="test", task_name="noop"
        )
        parent = self.playground.create_workflow(template, task_data={})
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test-step", display_name="Test Step"
        )
        workflow = wr.get_task()
        assert isinstance(workflow, Workflow)
        workflow.ensure_dynamic_data()
        assert workflow.dynamic_data is not None
        workflow.dynamic_data.parameter_summary = "summary"
        self.assertEqual(workflow.get_label(), "Test Step: summary")

    def test_get_label_workflow_template_name(self) -> None:
        template = self.playground.create_workflow_template(
            name="test", task_name="noop"
        )
        workflow = self.playground.create_workflow(template, task_data={})
        workflow.workflow_data_json["workflow_template_name"] = "template-name"
        workflow.save()

        self.assertEqual(
            workflow.get_task().get_label(),
            "template-name",
        )

    def test_get_label_workflow_include_parameter_summary(self) -> None:
        template = self.playground.create_workflow_template(
            name="test", task_name="noop"
        )
        workflow = self.playground.create_workflow(template, task_data={})
        workflow.dynamic_task_data = {"parameter_summary": "hello/amd64"}
        workflow.save()

        self.assertEqual(
            workflow.get_task().get_label(),
            "test: hello/amd64",
        )

    @preserve_db_task_registry()
    def test_get_description(self) -> None:
        task = self.make_concrete_workflow()
        self.assertEqual(task.get_description(), task.__doc__)

    @preserve_db_task_registry()
    def test_get_description_callback(self) -> None:
        def cb_function() -> None:
            """Test documentation."""
            raise NotImplementedError()

        callback = self.playground.create_internal_task(task_name="workflow")
        callback.workflow_data_json = {"step": "test"}
        task = self.make_concrete_workflow()
        task.work_request_callback = callback
        task.callback_test = cb_function  # type: ignore[attr-defined]
        assert task.__doc__ is not None
        assert cb_function.__doc__ is not None
        self.assertEqual(task.get_description(), cb_function.__doc__)

    @preserve_db_task_registry()
    def test_get_description_callback_no_step(self) -> None:
        callback = self.playground.create_internal_task(task_name="workflow")
        task = self.make_concrete_workflow()
        task.work_request_callback = callback
        self.assertEqual(task.get_description(), task.__doc__)

    @preserve_db_task_registry()
    def test_get_description_callback_invalid_step(self) -> None:
        callback = self.playground.create_internal_task(task_name="workflow")
        callback.workflow_data_json = {"step": "test"}
        task = self.make_concrete_workflow()
        task.work_request_callback = callback
        self.assertEqual(task.get_description(), task.__doc__)

    @preserve_db_task_registry()
    def test_get_description_callback_step_without_docstring(self) -> None:
        def cb_function() -> None:
            raise NotImplementedError()

        callback = self.playground.create_internal_task(task_name="workflow")
        callback.workflow_data_json = {"step": "test"}
        task = self.make_concrete_workflow()
        task.work_request_callback = callback
        task.callback_test = cb_function  # type: ignore[attr-defined]

        assert task.__doc__ is not None
        self.assertEqual(task.get_description(), task.__doc__)

    @preserve_db_task_registry()
    def test_compute_scheduler_tags_no_source_package_name(self) -> None:
        task = self.make_concrete_workflow()
        early_dynamic_task_data = mock.Mock()
        early_dynamic_task_data.get_source_package_name = mock.Mock(
            return_value=None
        )
        with mock.patch.object(
            task, "compute_dynamic_data", return_value=early_dynamic_task_data
        ):
            provided, _ = task.compute_scheduler_tags()
        self.assertCountEqual(
            provided.sets[ttags.ProvenanceProvided.USER], set()
        )

    @preserve_db_task_registry()
    def test_compute_scheduler_tags_with_source_package_name(self) -> None:
        task = self.make_concrete_workflow()
        early_dynamic_task_data = mock.Mock()
        early_dynamic_task_data.get_source_package_name = mock.Mock(
            return_value="foo"
        )
        with mock.patch.object(
            task, "compute_dynamic_data", return_value=early_dynamic_task_data
        ):
            provided, _ = task.compute_scheduler_tags()
        self.assertCountEqual(
            provided.sets[ttags.ProvenanceProvided.SYSTEM],
            {
                ttags.SOURCE_PACKAGE_PREFIX + "foo",
                ttags.SCOPE_PREFIX + "debusine",
                ttags.WORKSPACE_PREFIX + "debusine:System",
            },
        )


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

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

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.callback_test",
            create=True,
            return_value=True,
        ) as mock_noop_callback_test:
            self.assertTrue(orchestrate_workflow(wr))

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

    def test_workflow_callback_failure(self) -> None:
        """A workflow callback that returns False is marked as failed."""
        parent = self.playground.create_workflow(
            task_name="noop", mark_running=True
        )
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test"
        )
        wr.mark_pending()

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.callback_test",
            create=True,
            return_value=False,
        ) as mock_noop_callback_test:
            self.assertTrue(orchestrate_workflow(wr))

        mock_noop_callback_test.assert_called_once_with()
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.FAILURE)

    def test_workflow_callback_incremental_unterminated_dependencies(
        self,
    ) -> None:
        parent = self.playground.create_workflow(
            task_name="noop", mark_running=True
        )
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent,
            unblock_strategy=WorkRequest.UnblockStrategy.INCREMENTAL,
            step="test",
        )
        dep = parent.create_child_worker("noop")
        wr.add_dependency(dep)
        wr.mark_pending()

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.callback_test",
            create=True,
            return_value=True,
        ) as mock_noop_callback_test:
            self.assertTrue(orchestrate_workflow(wr))

        mock_noop_callback_test.assert_called_once_with()
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.BLOCKED)

    def test_workflow_callback_incremental_no_unterminated_dependencies(
        self,
    ) -> None:
        parent = self.playground.create_workflow(
            task_name="noop", mark_running=True
        )
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent,
            unblock_strategy=WorkRequest.UnblockStrategy.INCREMENTAL,
            step="test",
        )
        dep = parent.create_child_worker("noop")
        self.playground.advance_work_request(
            dep, result=WorkRequest.Results.SUCCESS
        )
        wr.add_dependency(dep)
        wr.mark_pending()

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.callback_test",
            create=True,
            return_value=True,
        ) as mock_noop_callback_test:
            self.assertTrue(orchestrate_workflow(wr))

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

    def test_workflow_callback_step_not_implemented(self) -> None:
        """A workflow callback must have a matching ``callback_*`` method."""
        parent = self.playground.create_workflow(
            task_name="noop", mark_running=True
        )
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test"
        )
        wr.mark_pending()
        expected_message = (
            "Orchestrator failed: Unhandled workflow callback step: test"
        )

        with (
            mock.patch(
                "debusine.server.workflows.noop.NoopWorkflow.callback_other",
                create=True,
                return_value=True,
            ) as mock_noop_callback_other,
            self.assertLogsContains(
                f"Error running work request Internal/workflow ({wr.id}): "
                f"{expected_message}",
                logger="debusine.server.workflows.base",
                level=logging.WARNING,
            ),
        ):
            self.assertFalse(orchestrate_workflow(wr))

        mock_noop_callback_other.assert_not_called()
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)
        self.assertEqual(
            wr.output_data,
            OutputData(
                errors=[
                    OutputDataError(
                        message=expected_message, code="orchestrator-failed"
                    )
                ]
            ),
        )

    @preserve_db_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(
            SampleWorkflow[BaseWorkflowData, ExampleDynamicData]
        ):
            TASK_NAME = "example"

            @override
            def compute_dynamic_data(self) -> ExampleDynamicData:
                return ExampleDynamicData(dynamic="foo")

            @override
            def populate(self) -> None:
                raise NotImplementedError()

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

        with mock.patch.object(ExampleWorkflow, "callback") as mock_callback:
            self.assertTrue(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_db_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 = "unset"

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

            @override
            def populate(self) -> None:
                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"
        )
        wr.mark_pending()

        with (
            mock.patch.object(
                ExampleWorkflow, "compute_dynamic_data"
            ) as mock_compute_dynamic_data,
            mock.patch.object(ExampleWorkflow, "callback") as mock_callback,
        ):
            self.assertTrue(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_internal_task(task_name="workflow")
        self.playground.advance_work_request(
            wr, mark_pending=True, use_data=False
        )
        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.assertFalse(orchestrate_workflow(wr))

        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)
        self.assertEqual(
            wr.output_data,
            OutputData(
                errors=[
                    OutputDataError(
                        message=expected_message, code="configure-failed"
                    )
                ]
            ),
        )

    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"
        )
        wr.mark_pending()
        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.assertFalse(orchestrate_workflow(wr))

        mock_noop_callback.assert_not_called()
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)
        assert isinstance(wr.output_data, OutputData)
        assert wr.output_data.errors is not None
        self.assertEqual(len(wr.output_data.errors), 1)
        self.assertTrue(
            wr.output_data.errors[0].message.startswith(expected_message)
        )
        self.assertEqual(wr.output_data.errors[0].code, "configure-failed")

    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"
        )
        wr.mark_pending()
        expected_message = (
            f"Failed to compute dynamic data for Workflow/noop ({parent.id}): "
            f"Boom"
        )

        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.assertLogsContains(
                f"Error running work request Internal/workflow ({wr.id}): "
                f"{expected_message}",
                logger="debusine.server.workflows.base",
                level=logging.WARNING,
            ),
        ):
            self.assertFalse(orchestrate_workflow(wr))

        mock_noop_callback.assert_not_called()
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)
        self.assertEqual(
            wr.output_data,
            OutputData(
                errors=[
                    OutputDataError(
                        message=expected_message, code="dynamic-data-failed"
                    )
                ]
            ),
        )

    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"
        )
        wr.mark_pending()
        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.assertFalse(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.ERROR)
        self.assertEqual(
            wr.output_data,
            OutputData(
                errors=[
                    OutputDataError(
                        message=expected_message, code="orchestrator-failed"
                    )
                ]
            ),
        )

    def test_workflow_callback_fails_without_context(self) -> None:
        """A callback failure by bare exception logs the exception name."""
        parent = self.playground.create_workflow(task_name="noop")
        wr = WorkRequest.objects.create_workflow_callback(
            parent=parent, step="test"
        )
        wr.mark_pending()
        expected_message = "Orchestrator failed: AssertionError()"

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

        self.assertEqual(
            wr.output_data,
            OutputData(
                errors=[
                    OutputDataError(
                        message=expected_message, code="orchestrator-failed"
                    )
                ]
            ),
        )

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

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

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.populate",
            side_effect=populate,
        ) as mock_noop_populate:
            self.assertTrue(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", mark_running=True
        )

        with mock.patch(
            "debusine.server.workflows.noop.NoopWorkflow.populate"
        ) as mock_noop_populate:
            self.assertTrue(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_db_task_registry()
    def test_workflow_empty_sub_workflow(self) -> None:
        """A workflow with an empty sub-workflow isn't prematurely completed."""

        class ExampleWorkflow(
            SampleWorkflow[BaseWorkflowData, BaseDynamicTaskData]
        ):
            TASK_NAME = "example"

            @override
            def populate(self) -> None:
                noop1 = self.work_request.create_child_workflow("noop")
                self.orchestrate_child(noop1)

                # Normally we'd orchestrate this sub-workflow as well, but
                # we skip that for the purposes of this test.
                self.work_request.create_child_workflow("noop")

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

        self.assertTrue(orchestrate_workflow(wr))

        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.RUNNING)
        self.assertQuerySetEqual(
            wr.children.order_by("id").values_list("status", "result"),
            [
                (WorkRequest.Statuses.COMPLETED, WorkRequest.Results.SUCCESS),
                (WorkRequest.Statuses.PENDING, WorkRequest.Results.NONE),
            ],
        )

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

        class ExampleDynamicData(BaseDynamicTaskData):
            dynamic: str

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

            @override
            def compute_dynamic_data(self) -> ExampleDynamicData:
                return ExampleDynamicData(dynamic="foo")

            @override
            def populate(self) -> None:
                self.work_request.create_child_worker("noop")

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

        self.assertTrue(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.configured_task_data = {"nonsense": ""}
        wr.save()
        child = wr.create_child_worker("noop")
        self.playground.advance_work_request(wr, mark_running=True)
        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.assertFalse(orchestrate_workflow(wr))

        mock_noop_populate.assert_not_called()
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)
        assert isinstance(wr.output_data, OutputData)
        assert wr.output_data.errors is not None
        self.assertEqual(len(wr.output_data.errors), 1)
        self.assertTrue(
            wr.output_data.errors[0].message.startswith(expected_message)
        )
        self.assertEqual(wr.output_data.errors[0].code, "configure-failed")

        child.refresh_from_db()
        self.assertEqual(child.status, WorkRequest.Statuses.ABORTED)
        assert isinstance(child.output_data, OutputData)
        assert child.output_data.errors is not None
        self.assertEqual(len(child.output_data.errors), 1)
        self.assertEqual(
            child.output_data.errors[0].message, "Error running workflow"
        )
        self.assertEqual(child.output_data.errors[0].code, "workflow-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", mark_running=True
        )
        child = wr.create_child_worker("noop")
        expected_message = "Failed to compute dynamic data: Boom"

        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.assertLogsContains(
                f"Error running work request Workflow/noop ({wr.id}): "
                f"{expected_message}",
                logger="debusine.server.workflows.base",
                level=logging.WARNING,
            ),
        ):
            self.assertFalse(orchestrate_workflow(wr))

        mock_noop_populate.assert_not_called()
        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)
        self.assertEqual(
            wr.output_data,
            OutputData(
                errors=[
                    OutputDataError(
                        message=expected_message, code="dynamic-data-failed"
                    )
                ]
            ),
        )

        child.refresh_from_db()
        self.assertEqual(child.status, WorkRequest.Statuses.ABORTED)
        assert isinstance(child.output_data, OutputData)
        assert child.output_data.errors is not None
        self.assertEqual(len(child.output_data.errors), 1)
        self.assertEqual(
            child.output_data.errors[0].message, "Error running workflow"
        )
        self.assertEqual(child.output_data.errors[0].code, "workflow-error")

    def test_workflow_wrong_status(self) -> None:
        """A workflow in the wrong status is skipped."""
        wr = self.playground.create_workflow(
            task_name="noop", mark_pending=False
        )
        expected_message = "Failed to mark work request as running"

        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.assertFalse(orchestrate_workflow(wr))

        mock_noop_populate.assert_not_called()
        wr.refresh_from_db()
        # The work request remains blocked.
        self.assertEqual(wr.status, WorkRequest.Statuses.BLOCKED)
        # No errors are logged in output_data, because the work request may
        # make progress in the future.
        self.assertIsNone(wr.output_data)

    def test_workflow_fails(self) -> None:
        """A workflow that fails is logged."""
        wr = self.playground.create_workflow(
            task_name="noop", mark_running=True
        )
        child = wr.create_child_worker("noop")
        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.assertFalse(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.ERROR)
        self.assertEqual(
            wr.output_data,
            OutputData(
                errors=[
                    OutputDataError(
                        message=expected_message, code="orchestrator-failed"
                    )
                ]
            ),
        )

        child.refresh_from_db()
        self.assertEqual(child.status, WorkRequest.Statuses.ABORTED)
        assert isinstance(child.output_data, OutputData)
        assert child.output_data.errors is not None
        self.assertEqual(len(child.output_data.errors), 1)
        self.assertEqual(
            child.output_data.errors[0].message, "Error running workflow"
        )
        self.assertEqual(child.output_data.errors[0].code, "workflow-error")

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

        def populate() -> None:
            children = [wr.create_child_worker("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:
            self.assertTrue(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)

    @preserve_db_task_registry()
    def test_workflow_caches_lookups_for_children(self) -> None:
        """
        A workflow can cache lookups for its children.

        Any additional lookups performed by the children aren't cached
        outside their subtree.
        """
        artifacts = [self.playground.create_artifact()[0] for _ in range(2)]
        artifact_lookups = [
            f"{artifact.id}@artifacts" for artifact in artifacts
        ]
        results = {
            lookup: LookupResultArtifact(lookup=lookup, artifact=artifact)
            for artifact, lookup in zip(artifacts, artifact_lookups)
        }

        class CachingSubWorkflow(
            SampleWorkflow[BaseWorkflowData, BaseDynamicTaskData]
        ):
            TASK_NAME = "caching_sub"

            @override
            def populate(self_) -> None:
                with self.assertNumQueries(0):
                    self.assertEqual(
                        self_.work_request.lookup_single(artifact_lookups[0]),
                        results[artifact_lookups[0]],
                    )
                self_.work_request.lookup_single(
                    artifact_lookups[1],
                    precache_result=results[artifact_lookups[1]],
                )

        class CachingRootWorkflow(
            SampleWorkflow[BaseWorkflowData, BaseDynamicTaskData]
        ):
            TASK_NAME = "caching_root"

            @override
            def populate(self_) -> None:
                self_.work_request.lookup_single(
                    artifact_lookups[0],
                    precache_result=results[artifact_lookups[0]],
                )
                wr, _ = self_.work_request_ensure_child_workflow(
                    task_name="caching_sub",
                    workflow_data=WorkRequestWorkflowData(
                        display_name="sub", step="sub"
                    ),
                )
                self_.orchestrate_child(wr)
                with self.assertNumQueries(1):
                    self.assertEqual(
                        self_.work_request.lookup_single(artifact_lookups[1]),
                        results[artifact_lookups[1]],
                    )

        wr = self.playground.create_workflow(
            task_name="caching_root", mark_running=True
        )

        self.assertTrue(orchestrate_workflow(wr))

        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.SUCCESS)
        self.assertQuerySetEqual(
            wr.children.order_by("id").values_list("status", "result"),
            [(WorkRequest.Statuses.COMPLETED, WorkRequest.Results.SUCCESS)],
        )

    def test_wrong_task_type(self) -> None:
        """Attempts to orchestrate non-workflows are logged."""
        wr = self.playground.create_worker_task(
            mark_pending=True, assign_new_worker=True
        )
        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.assertFalse(orchestrate_workflow(wr))

        wr.refresh_from_db()
        self.assertEqual(wr.status, WorkRequest.Statuses.COMPLETED)
        self.assertEqual(wr.result, WorkRequest.Results.ERROR)
        self.assertEqual(
            wr.output_data,
            OutputData(
                errors=[
                    OutputDataError(
                        message=expected_message, code="configure-failed"
                    )
                ]
            ),
        )


class TaskConfigurationTests(TestCase):
    """test :py:class:`Workflow` task configuration handling."""

    scenario = scenarios.DefaultContext()

    workflow_wr: ClassVar[WorkRequest]

    @override
    @classmethod
    def setUpTestData(cls) -> None:
        super().setUpTestData()
        cls.workflow_wr = cls.playground.create_workflow()

    @override
    def setUp(self) -> None:
        super().setUp()
        workflow = self.workflow_wr.get_task()
        assert isinstance(workflow, Workflow)
        self.workflow = workflow

    def test_work_request_ensure_child(self) -> None:
        package = self.playground.create_minimal_binary_package_artifact()
        for task_type, task_name, workflow_data, kwargs in (
            ("worker", "noop", WorkRequestWorkflowData(step="test"), {}),
            ("server", "servernoop", WorkRequestWorkflowData(step="test"), {}),
            ("internal", "workflow", WorkRequestWorkflowData(step="test"), {}),
            ("workflow", "noop", WorkRequestWorkflowData(step="test"), {}),
            ("signing", "noop", WorkRequestWorkflowData(step="test"), {}),
            (
                "wait",
                "externaldebsign",
                WorkRequestWorkflowData(step="test", needs_input=True),
                {"task_data": ExternalDebsignData(unsigned=package.pk)},
            ),
        ):
            with self.subTest(task_type=task_type, task_name=task_name):
                factory = partial(
                    getattr(
                        self.workflow, f"work_request_ensure_child_{task_type}"
                    ),
                    task_name=task_name,
                    workflow_data=workflow_data,
                    **kwargs,
                )
                wr, created = factory()
                self.assertIsNone(wr.configured_task_data)
                self.assertEqual(wr.status, WorkRequest.Statuses.BLOCKED)
                self.assertTrue(created)

                self.assertEqual(factory(), (wr, False))

                if task_type in {"worker", "internal", "workflow", "signing"}:
                    wr.mark_pending()
                    wr.mark_aborted()
                    with context.disable_permission_checks():
                        retried_wr = wr.retry()

                    self.assertEqual(factory(), (retried_wr, False))


class RTExampleWorkflowData(BaseWorkflowData):
    """Example Workflow Data for RuntimeParametersTests."""

    dynamic: str
    nested: dict[str, str]
    multiple: list[str]


class RuntimeParametersTests(TestCase):
    """Test Workflow.validate_runtime_parameters()."""

    workflow: ClassVar[
        type[Workflow[RTExampleWorkflowData, BaseDynamicTaskData]]
    ]

    @override
    @classmethod
    @preserve_db_task_registry()
    def setUpTestData(cls) -> None:
        """Add a sample Workflow to the test fixture."""
        super().setUpTestData()

        class RTExampleWorkflow(
            SampleWorkflow[RTExampleWorkflowData, BaseDynamicTaskData]
        ):
            TASK_NAME = "example"

        # https://github.com/python/mypy/issues/4717
        cls.workflow = RTExampleWorkflow  # type: ignore[type-abstract]

    def test_validate_any(self) -> None:
        self.workflow.validate_runtime_parameters(RuntimeParameter.ANY)

    def test_validate_empty_dict(self) -> None:
        self.workflow.validate_runtime_parameters({})

    def test_validate_parameter_any(self) -> None:
        self.workflow.validate_runtime_parameters(
            {"dynamic": RuntimeParameter.ANY}
        )

    def test_validate_parameter_values(self) -> None:
        self.workflow.validate_runtime_parameters({"dynamic": ["foo", "bar"]})

    def test_validate_parameter_values_wrong_type(self) -> None:
        with self.assertRaisesRegex(
            WorkflowValidationError,
            r"dynamic is set to \[1\], not a list of valid values\.",
        ):
            self.workflow.validate_runtime_parameters({"dynamic": [1]})

    def test_validate_parameter_single_value(self) -> None:
        with self.assertRaisesRegex(
            WorkflowValidationError,
            (
                r"dynamic is set to a string \('foo'\), not a list of valid "
                r"values\."
            ),
        ):
            self.workflow.validate_runtime_parameters({"dynamic": "foo"})

    def test_validate_unknown_parameter(self) -> None:
        with self.assertRaisesRegex(
            WorkflowValidationError,
            (
                r"unknown_parameter is not a task_data parameter for "
                r"RTExampleWorkflow"
            ),
        ):
            self.workflow.validate_runtime_parameters(
                {"unknown_parameter": RuntimeParameter.ANY}
            )
