comet-ml--opik/sdks/python/tests/unit/cli/_migrate_prompt_helpers.py

"""Shared mock helpers for ``opik migrate prompt`` test modules.

Used by ``test_migrate_prompt_planner.py``,
``test_migrate_prompt_executor.py``, and
``test_migrate_prompt_version_replay.py``. Lives as a plain module (not
a conftest.py) because these are helper classes, not pytest fixtures.
"""

from __future__ import annotations

from typing import Any, List, Optional
from unittest.mock import MagicMock


class _PromptRow:
    """Minimal stand-in for ``PromptPublic`` rows returned by ``get_prompts``."""

    def __init__(
        self,
        id: str,
        name: str,
        *,
        description: Optional[str] = None,
        tags: Optional[List[str]] = None,
        # None = workspace-scoped (legacy v1 prompts that never had a
        # project). Tests that want a project-scoped source pass an id.
        project_id: Optional[str] = None,
        template_structure: Optional[str] = "text",
    ) -> None:
        self.id = id
        self.name = name
        self.description = description
        self.tags = tags
        self.project_id = project_id
        self.template_structure = template_structure


class _PromptVersionRow:
    """Stand-in for ``PromptVersionPublic`` rows from ``get_prompt_versions``.

    Field set mirrors the Fern model: ``id``, ``prompt_id``, ``commit``,
    ``template`` (required), plus optional metadata / type / change_description
    / tags / environments / template_structure.
    """

    def __init__(
        self,
        id: str,
        prompt_id: str,
        commit: str,
        template: str,
        *,
        metadata: Optional[Any] = None,
        type: Optional[str] = None,
        change_description: Optional[str] = None,
        tags: Optional[List[str]] = None,
        environments: Optional[List[str]] = None,
        template_structure: Optional[str] = "text",
    ) -> None:
        self.id = id
        self.prompt_id = prompt_id
        self.commit = commit
        self.template = template
        self.metadata = metadata
        self.type = type
        self.change_description = change_description
        self.tags = tags
        self.environments = environments
        self.template_structure = template_structure


class _Page:
    def __init__(self, content: List[Any]) -> None:
        self.content = content


def _named(name: str) -> MagicMock:
    obj = MagicMock()
    obj.name = name
    return obj


def _planner_rest_client(
    find_prompts_side_effects: List[_Page],
    *,
    target_project_exists: bool = True,
    workspace_project_names: Optional[List[str]] = None,
) -> MagicMock:
    """Build a rest_client mock for direct prompt-planner unit tests.

    ``find_prompts_side_effects`` is the side-effect list for
    ``rest_client.prompts.get_prompts``; the planner calls it twice (once
    for source resolution, once for the rename-collision preflight) so
    callers pass two pages.
    """
    rest_client = MagicMock()
    if target_project_exists:
        target_project = MagicMock()
        target_project.id = "target-project-id"
        rest_client.projects.retrieve_project.return_value = target_project
    else:
        from opik.rest_api.core.api_error import ApiError

        rest_client.projects.retrieve_project.side_effect = ApiError(
            status_code=404, body={}
        )
        candidates = [_named(n) for n in (workspace_project_names or [])]
        rest_client.projects.find_projects.return_value = _Page(candidates)
    rest_client.prompts.get_prompts.side_effect = find_prompts_side_effects
    return rest_client


def _planner_client(rest_client: MagicMock) -> MagicMock:
    """Wrap a rest_client mock as an ``opik.Opik``-shaped client.

    The resolver routes project-id lookups through ``client.get_project``
    (high-level surface) so the project-name resolution still benefits
    from any test-side stubs on ``rest_client.projects.get_project_by_id``.
    """
    client = MagicMock()
    client.rest_client = rest_client
    client.get_project = MagicMock(
        side_effect=lambda id: rest_client.projects.get_project_by_id(id=id)
    )
    return client