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

"""DB-aware mock field resolver."""

from debusine.db.models import Artifact, Collection
from debusine.db.tasks.models import DBArtifactInfo, DBInputCollectionSingle
from debusine.tasks.inputs.mock_resolver import (
    MockFieldResolver as BaseMockFieldResolver,
)


class MockFieldResolver(BaseMockFieldResolver):
    """DB-aware mock field resolver."""

    def set_artifact(self, name: str, artifact: Artifact) -> None:
        """Set the value of an input field for an artifact lookup."""
        self.field_values[name] = DBArtifactInfo.from_artifact(
            artifact, label=name, lookup=artifact.id
        )

    def set_collection(self, name: str, collection: Collection) -> None:
        """Set the value of an input field for a collection lookup."""
        self.field_values[name] = DBInputCollectionSingle.from_collection(
            collection, label=name, lookup=collection.id
        )
