chore: import upstream snapshot with attribution

sdks/python/tests/unit/test_cli_changes.py

@@ -0,0 +1,508 @@
+"""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