comet-ml--opik/sdks/python/tests/unit/test_cli_changes.py

"""Unit tests for CLI changes: build_import_metadata, _validate_include,
prompt type case-insensitivity, export_traces unlimited pagination,
and 'all' subcommand registration."""

import json
import tempfile
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest
import click
from click.testing import CliRunner

from opik.cli import cli
from opik.cli.imports.utils import (
    build_import_metadata,
    _TRACE_IMPORT_FIELDS,
    _SPAN_IMPORT_FIELDS,
    _EXPERIMENT_IMPORT_FIELDS,
)
from opik.api_objects.prompt.types import PromptType


# ---------------------------------------------------------------------------
# build_import_metadata
# ---------------------------------------------------------------------------


class TestBuildImportMetadata:
    def test_build_import_metadata__no_matching_fields_no_existing_metadata__returns_none(
        self,
    ):
        source = {"name": "my-trace"}
        result = build_import_metadata(source, _TRACE_IMPORT_FIELDS, None)
        assert result is None

    def test_build_import_metadata__no_matching_fields_with_existing__returns_existing_unchanged(
        self,
    ):
        source = {"name": "my-trace"}
        existing = {"key": "value"}
        result = build_import_metadata(source, _TRACE_IMPORT_FIELDS, existing)
        assert result is existing

    def test_build_import_metadata__fields_present_no_existing__happyflow(self):
        source = {"created_by": "alice", "created_at": "2024-01-01T00:00:00Z"}
        result = build_import_metadata(source, ["created_by", "created_at"], None)
        assert result == {
            "_import_created_by": "alice",
            "_import_created_at": "2024-01-01T00:00:00Z",
        }

    def test_build_import_metadata__fields_present_merged_with_existing__happyflow(
        self,
    ):
        source = {"created_by": "bob", "ttft": 0.5}
        existing = {"custom_key": "custom_val"}
        result = build_import_metadata(source, ["created_by", "ttft"], existing)
        assert result == {
            "custom_key": "custom_val",
            "_import_created_by": "bob",
            "_import_ttft": 0.5,
        }

    def test_build_import_metadata__none_field_values__are_skipped(self):
        source = {"created_by": None, "last_updated_by": "carol"}
        result = build_import_metadata(source, ["created_by", "last_updated_by"], None)
        assert result == {"_import_last_updated_by": "carol"}

    def test_build_import_metadata__all_none_values_with_existing__returns_existing(
        self,
    ):
        source = {"created_by": None, "created_at": None}
        existing = {"x": 1}
        result = build_import_metadata(source, ["created_by", "created_at"], existing)
        assert result is existing

    def test_build_import_metadata__with_existing__does_not_mutate_existing(self):
        source = {"created_by": "dave"}
        existing = {"orig": "val"}
        build_import_metadata(source, ["created_by"], existing)
        assert existing == {"orig": "val"}

    def test_build_import_metadata__span_fields_subset__happyflow(self):
        source = {
            "created_at": "2024-01-01",
            "created_by": "user",
            "last_updated_at": "2024-01-02",
            "last_updated_by": "user2",
            "ttft": 1.2,
        }
        result = build_import_metadata(source, _SPAN_IMPORT_FIELDS, None)
        assert set(result.keys()) == {
            "_import_created_at",
            "_import_created_by",
            "_import_last_updated_at",
            "_import_last_updated_by",
            "_import_ttft",
        }

    def test_build_import_metadata__experiment_fields_subset__happyflow(self):
        source = {
            "created_at": "2024-01-01",
            "created_by": "user",
            "last_updated_at": "2024-01-02",
            "last_updated_by": "user2",
        }
        result = build_import_metadata(source, _EXPERIMENT_IMPORT_FIELDS, None)
        assert len(result) == 4
        assert "_import_created_by" in result


# ---------------------------------------------------------------------------
# _validate_include (export side)
# ---------------------------------------------------------------------------


class TestValidateIncludeExport:
    """Tests for _validate_include callback in exports/all.py."""

    def test_valid_include_option_accepted(self):
        runner = CliRunner()
        result = runner.invoke(
            cli,
            [
                "export",
                "default",
                "proj",
                "all",
                "--include",
                "datasets,prompts",
                "--help",
            ],
        )
        # --help always exits 0 regardless of option values
        assert result.exit_code == 0

    def test_invalid_include_option_raises_error(self):
        runner = CliRunner()
        with patch("opik.cli.exports.all.opik.Opik"):
            result = runner.invoke(
                cli,
                ["export", "default", "proj", "all", "--include", "invalid_type"],
            )
        assert result.exit_code != 0
        assert "Invalid" in result.output or "invalid" in result.output.lower()

    def test_case_insensitive_include(self):
        """_validate_include normalises to lower-case; upper-case valid values pass."""
        from opik.cli.exports.all import _validate_include

        ctx = MagicMock(spec=click.Context)
        param = MagicMock(spec=click.Parameter)
        result = _validate_include(ctx, param, "Datasets,PROMPTS")
        assert result == ["datasets", "prompts"]

    def test_all_valid_types_accepted(self):
        from opik.cli.exports.all import _validate_include

        ctx = MagicMock(spec=click.Context)
        param = MagicMock(spec=click.Parameter)
        result = _validate_include(ctx, param, "datasets,prompts,traces,experiments")
        assert set(result) == {"datasets", "prompts", "traces", "experiments"}

    def test_invalid_type_raises_bad_parameter(self):
        from opik.cli.exports.all import _validate_include

        ctx = MagicMock(spec=click.Context)
        param = MagicMock(spec=click.Parameter)
        with pytest.raises(click.BadParameter, match="Invalid"):
            _validate_include(ctx, param, "datasets,unknown")


# ---------------------------------------------------------------------------
# _validate_include (import side)
# ---------------------------------------------------------------------------


class TestValidateIncludeImport:
    """Tests for _validate_include callback in imports/all.py."""

    def test_valid_include_types(self):
        from opik.cli.imports.all import _validate_include

        ctx = MagicMock(spec=click.Context)
        param = MagicMock(spec=click.Parameter)
        result = _validate_include(ctx, param, "datasets,experiments")
        assert result == ["datasets", "experiments"]

    def test_invalid_include_raises_bad_parameter(self):
        from opik.cli.imports.all import _validate_include

        ctx = MagicMock(spec=click.Context)
        param = MagicMock(spec=click.Parameter)
        with pytest.raises(click.BadParameter, match="Invalid"):
            _validate_include(ctx, param, "projects")

    def test_empty_segments_ignored(self):
        from opik.cli.imports.all import _validate_include

        ctx = MagicMock(spec=click.Context)
        param = MagicMock(spec=click.Parameter)
        # Leading/trailing commas and spaces
        result = _validate_include(ctx, param, " datasets , , prompts ")
        assert result == ["datasets", "prompts"]


# ---------------------------------------------------------------------------
# Prompt type case-insensitivity fix
# ---------------------------------------------------------------------------


class TestPromptTypeResolution:
    """Verify that the import correctly handles UPPERCASE prompt type strings."""

    def test_uppercase_mustache_resolves(self):
        assert PromptType("MUSTACHE".lower()) == PromptType.MUSTACHE

    def test_uppercase_jinja2_resolves(self):
        assert PromptType("JINJA2".lower()) == PromptType.JINJA2

    def test_lowercase_mustache_resolves(self):
        assert PromptType("mustache") == PromptType.MUSTACHE

    def test_import_prompts_handles_uppercase_type(self):
        """End-to-end: import_prompts_from_directory resolves 'MUSTACHE' type."""
        from opik.cli.imports.prompt import import_prompts_from_directory

        prompt_data = {
            "name": "test-prompt",
            "current_version": {
                "prompt": "Hello {{name}}",
                "type": "MUSTACHE",
                "template_structure": "text",
                "metadata": None,
            },
        }

        mock_client = MagicMock()
        mock_prompt = MagicMock()
        mock_client.create_prompt.return_value = mock_prompt

        with tempfile.TemporaryDirectory() as tmp:
            prompt_file = Path(tmp) / "prompt_test.json"
            prompt_file.write_text(json.dumps(prompt_data))

            result = import_prompts_from_directory(
                client=mock_client,
                source_dir=Path(tmp),
                project_name="test-project",
                dry_run=False,
                name_pattern=None,
                debug=False,
            )

            # Should have imported one prompt, not skipped it
            assert result.get("prompts", 0) == 1
            assert result.get("prompts_skipped", 0) == 0

    def test_import_prompts_handles_unknown_type_falls_back_to_mustache(self):
        """Unknown prompt type falls back to MUSTACHE."""
        from opik.cli.imports.prompt import import_prompts_from_directory

        prompt_data = {
            "name": "test-prompt2",
            "current_version": {
                "prompt": "Hello",
                "type": "COMPLETELY_UNKNOWN_TYPE",
                "template_structure": "text",
                "metadata": None,
            },
        }

        mock_client = MagicMock()

        with tempfile.TemporaryDirectory() as tmp:
            prompt_file = Path(tmp) / "prompt_test2.json"
            prompt_file.write_text(json.dumps(prompt_data))

            result = import_prompts_from_directory(
                client=mock_client,
                source_dir=Path(tmp),
                project_name="test-project",
                dry_run=False,
                name_pattern=None,
                debug=False,
            )

            # Falls back to MUSTACHE → still imports successfully
            assert result.get("prompts", 0) == 1


# ---------------------------------------------------------------------------
# export_traces: max_results=None fetches all pages
# ---------------------------------------------------------------------------


class TestExportTracesMaxResultsNone:
    """Verify that passing max_results=None to export_traces exhausts all pages."""

    def _make_mock_trace(self, trace_id: str) -> MagicMock:
        t = MagicMock()
        t.id = trace_id
        t.name = f"trace-{trace_id}"
        t.model_dump.return_value = {
            "id": trace_id,
            "name": f"trace-{trace_id}",
            "start_time": None,
            "end_time": None,
            "input": {},
            "output": {},
            "metadata": {},
            "tags": [],
            "feedback_scores": [],
            "error_info": None,
            "thread_id": None,
            "created_at": None,
            "created_by": None,
            "last_updated_at": None,
            "last_updated_by": None,
            "visibility_mode": None,
            "ttft": None,
            "project_name": "test-project",
        }
        return t

    def _make_page(self, traces):
        page = MagicMock()
        page.content = traces
        return page

    def test_all_traces_exported_when_max_results_is_none(self):
        """When max_results=None every trace on the page is exported (no early stop)."""
        from opik.cli.exports.project import export_traces

        traces = [self._make_mock_trace(f"t{i}") for i in range(3)]

        mock_client = MagicMock()
        # Single page of 3 traces (< page_size=100, so loop exits after this page)
        mock_client.rest_client.traces.get_traces_by_project.return_value = (
            self._make_page(traces)
        )
        mock_client.search_spans.return_value = []

        with tempfile.TemporaryDirectory() as tmp:
            exported, skipped, had_errors = export_traces(
                client=mock_client,
                project_name="test-project",
                project_dir=Path(tmp),
                max_results=None,
                filter_string=None,
            )

        assert exported == 3
        assert skipped == 0
        assert had_errors is False

    def test_max_results_limits_export(self):
        """When max_results=1 only 1 trace is requested from the API."""
        from opik.cli.exports.project import export_traces

        # With max_results=1, current_page_size=min(100,1)=1, so API returns at most 1
        traces = [self._make_mock_trace("t1")]

        mock_client = MagicMock()
        mock_client.rest_client.traces.get_traces_by_project.return_value = (
            self._make_page(traces)
        )
        mock_client.search_spans.return_value = []

        with tempfile.TemporaryDirectory() as tmp:
            exported, skipped, had_errors = export_traces(
                client=mock_client,
                project_name="test-project",
                project_dir=Path(tmp),
                max_results=1,
                filter_string=None,
            )

        assert exported == 1
        # Verify the API was called with size=1 (respects max_results)
        call_kwargs = mock_client.rest_client.traces.get_traces_by_project.call_args[1]
        assert call_kwargs["size"] == 1


# ---------------------------------------------------------------------------
# CLI 'all' subcommand registration
# ---------------------------------------------------------------------------


class TestAllCommandRegistered:
    def test_export_all_help_is_accessible(self):
        runner = CliRunner()
        result = runner.invoke(cli, ["export", "default", "proj", "all", "--help"])
        assert result.exit_code == 0
        assert "all" in result.output.lower()
        assert "--include" in result.output

    def test_import_all_help_is_accessible(self):
        runner = CliRunner()
        result = runner.invoke(cli, ["import", "default", "proj", "all", "--help"])
        assert result.exit_code == 0
        assert "all" in result.output.lower()
        assert "--include" in result.output

    def test_export_group_help_lists_all(self):
        runner = CliRunner()
        result = runner.invoke(cli, ["export", "--help"])
        assert result.exit_code == 0
        assert "all" in result.output

    def test_import_group_help_lists_all(self):
        runner = CliRunner()
        result = runner.invoke(cli, ["import", "--help"])
        assert result.exit_code == 0
        assert "all" in result.output

    def test_export_missing_subcommand_error_mentions_all(self):
        """When no subcommand is given the error message should list 'all'."""
        runner = CliRunner()
        result = runner.invoke(cli, ["export", "default", "proj"])
        # Non-zero exit or the error message includes "all"
        assert "all" in result.output


# ---------------------------------------------------------------------------
# _export_all_experiments: semaphore callback prevents deadlock
# ---------------------------------------------------------------------------


class TestExportAllExperimentsSemaphore:
    """Verify the semaphore done_callback in _export_all_experiments prevents deadlock."""

    def test_semaphore_callback_prevents_deadlock__more_than_capacity_experiments__all_exported(
        self,
    ):
        """Submitting N > max_workers*2 experiments must complete without deadlock.

        With max_workers=2 the semaphore capacity is 4.  Submitting 6 experiments
        would block the submission loop forever (the as_completed drain never
        starts) unless the done_callback on each future releases the semaphore.
        """
        from opik.cli.exports.all import _export_all_experiments
        from types import SimpleNamespace

        num_experiments = 6
        max_workers = 2  # semaphore capacity = max_workers * 2 = 4; 6 > 4

        experiments = [
            SimpleNamespace(id=f"exp-{i}", name=f"experiment-{i}")
            for i in range(num_experiments)
        ]

        def fake_export_by_id(*args, **kwargs):
            return (
                {
                    "datasets": 0,
                    "datasets_skipped": 0,
                    "prompts": 0,
                    "prompts_skipped": 0,
                    "traces": 0,
                    "traces_skipped": 0,
                },
                1,
                None,
            )

        with tempfile.TemporaryDirectory() as tmp:
            workspace_root = Path(tmp)
            experiments_dir = workspace_root / "experiments"
            experiments_dir.mkdir()

            with (
                patch(
                    "opik.cli.exports.all._paginate_experiments",
                    return_value=iter(experiments),
                ),
                patch(
                    "opik.cli.exports.all.export_experiment_by_id",
                    side_effect=fake_export_by_id,
                ),
                patch(
                    "opik.cli.exports.all.export_collected_trace_ids",
                    return_value=(0, 0),
                ),
            ):
                (
                    exp_exported,
                    exp_skipped,
                    traces_exported,
                    traces_skipped,
                    _had_errors,
                ) = _export_all_experiments(
                    client=MagicMock(),
                    project_dir=workspace_root,
                    project_name="proj",
                    project_id=None,
                    experiments_dir=experiments_dir,
                    max_results=None,
                    force=False,
                    debug=False,
                    format="json",
                    max_workers=max_workers,
                )

        # If the semaphore callback was missing the test would hang before this line.
        assert exp_exported + exp_skipped == num_experiments
        assert exp_exported == num_experiments