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

"""Common test-helper code involving wait tasks."""

from typing import Any, override

import debusine.tasks.tags as ttags
import debusine.worker.tags as wtags
from debusine.db.models import WorkRequest
from debusine.server.tasks.tests.helpers import DBTaskTestCase
from debusine.server.tasks.wait import BaseWaitTask
from debusine.server.workflows.models import WorkRequestWorkflowData
from debusine.test.django import TestCase


class WaitTaskTestCase[TaskClass: BaseWaitTask[Any, Any]](
    DBTaskTestCase[TaskClass], TestCase
):
    """Common base for wait task tests."""

    #: Set to true to mark the task as running on creation
    mark_running: bool

    @override
    def setUp(self) -> None:
        super().setUp()
        self.mark_running = False

    @override
    def work_request_create(self, **kwargs: Any) -> WorkRequest:
        task_data = self.task_data.copy()
        task_data.update(kwargs)
        work_request = self.playground.create_wait_task(
            task_name=self.task_class.name,
            task_data=self.task_class.task_data_type(**task_data),
            workflow_data=WorkRequestWorkflowData(needs_input=True),
        )
        if self.mark_running:
            self.playground.advance_work_request(
                work_request, mark_running=True
            )
        return work_request

    @override
    def build_expected_tag_candidates(
        self, **kwargs: set[str]
    ) -> tuple[ttags.TaskProvidedTags, wtags.TaskRequiredTags]:
        provided, required = super().build_expected_tag_candidates(**kwargs)
        required.add(
            wtags.ProvenanceRequired.SYSTEM, {wtags.WORKER_TYPE_NOT_ASSIGNABLE}
        )
        return provided, required
