"""End-to-end tests for ``opik migrate dataset`` against a real Opik backend. Covers the plain-dataset path: full version replay, plus the experiment + trace + span cascade that rides along with the dataset. The test-suite path lives in ``test_migrate_test_suite_e2e.py``. Each test: 1. Seeds a multi-version source dataset directly via the REST API (so we control every per-version delta the migration has to replay) and, where the test calls for it, an experiment + traces + spans attached to one of the source versions 2. Runs ``opik migrate dataset`` as a subprocess so the actual CLI entrypoint, Click group, and exit-code handling are exercised 3. Reads back the target via the raw REST stream + wire type (the SDK helper drops per-item tags) and asserts per-version content + display-order fidelity, plus -- where relevant -- destination experiment + trace + span fidelity and FK remapping Shared helpers live in ``conftest.py``. """ from __future__ import annotations import json from pathlib import Path from typing import Iterator import pytest import opik from opik import synchronization from ...conftest import random_chars from ...testlib import generate_project_name from ._cascade_comparison import compare_cascade from .conftest import ( apply_changes, chronological_versions, create_dataset_shell, destination_experiment_items, destination_feedback_scores_for_trace, destination_spans_for_trace, display_order, find_destination_experiment, item_hashes, run_migrate_cli, seed_experiment_with_trace_tree, stream_items_wire, ) # Per ``sdks/python/AGENTS.md``: every e2e module sources PROJECT_NAME from # ``generate_project_name("e2e", __name__)`` so backend project names are # isolated per test module + the autouse ``configure_e2e_tests_env`` fixture # can patch ``OPIK_PROJECT_NAME`` to match. PROJECT_NAME = generate_project_name("e2e", __name__) @pytest.fixture def dataset_name() -> Iterator[str]: yield f"e2e-migrate-{random_chars()}" class TestMigrateDatasetVersionReplay: """Default ``opik migrate dataset`` flow: full version-history replay. Pin the slice 2 contract: target version count == source version count, per-version content set-equal under hash, display order preserved at every version. """ def test_three_version_dataset_with_mixed_deltas_round_trips( self, opik_client: opik.Opik, source_project_name: str, target_project_name: str, dataset_name: str, tmp_path: Path, ) -> None: rest = opik_client.rest_client # Seed: 3 versions exercising every delta kind. # v1: adds Q1, Q2, Q3. # v2: edits Q1 (data change), adds Q4. # v3: edits Q3 (data change), deletes Q2, adds Q5. source_id = create_dataset_shell(rest, dataset_name, source_project_name) # v1 of a plain dataset is created via the REST insert path (mirrors # what ``Dataset.insert`` does at higher level). Use create_or_update # so we get exactly one BE version with all three items. from opik import id_helpers from opik.rest_api.types.dataset_item_write import DatasetItemWrite rest.datasets.create_or_update_dataset_items( dataset_id=source_id, items=[ DatasetItemWrite(source="manual", data={"q": "Q1", "a": "A1"}), DatasetItemWrite(source="manual", data={"q": "Q2", "a": "A2"}), DatasetItemWrite(source="manual", data={"q": "Q3", "a": "A3"}), ], batch_group_id=id_helpers.generate_id(), ) v1 = rest.datasets.list_dataset_versions(id=source_id, page=1, size=1).content[ 0 ] v1_items = stream_items_wire( rest, dataset_name=dataset_name, project_name=source_project_name, version_hash=v1.version_hash, ) by_q = {item.data["q"]: item for item in v1_items if item.data} # v2: edit Q1's answer + add Q4. v2_id = apply_changes( rest, source_id, base_version_id=v1.id, edited_items=[ {"id": by_q["Q1"].id, "data": {"q": "Q1", "a": "A1-EDITED"}}, ], added_items=[{"data": {"q": "Q4", "a": "A4"}, "source": "manual"}], change_description="v2 — edit Q1, add Q4", ) # v3: edit Q3, delete Q2, add Q5. apply_changes( rest, source_id, base_version_id=v2_id, edited_items=[ {"id": by_q["Q3"].id, "data": {"q": "Q3", "a": "A3-EDITED"}}, ], deleted_ids=[by_q["Q2"].id], added_items=[{"data": {"q": "Q5", "a": "A5"}, "source": "manual"}], change_description="v3 — delete Q2, edit Q3, add Q5", ) # Snapshot source expectations per version. src_versions = chronological_versions(rest, source_id) assert len(src_versions) == 3 expected_hashes = [] expected_orders = [] for v in src_versions: items = stream_items_wire( rest, dataset_name=dataset_name, project_name=source_project_name, version_hash=v.version_hash, ) expected_hashes.append(item_hashes(items)) expected_orders.append(display_order(items)) # ── Seed an experiment on v1 items so the cascade has something to # round-trip. Regular-dataset experiments carry per-trace feedback # scores (test suites carry assertion_results -- covered in # test_migrate_test_suite_e2e.py). Each item gets one trace with a # root + 1 LLM child span and a feedback score on the trace. ── experiment_name = f"e2e-exp-{random_chars()}" v1_item_ids = [by_q["Q1"].id, by_q["Q2"].id, by_q["Q3"].id] cascade_seed = seed_experiment_with_trace_tree( rest, experiment_name=experiment_name, dataset_name=dataset_name, dataset_id=source_id, dataset_version_id=v1.id, project_name=source_project_name, item_ids=v1_item_ids, experiment_config={"runner": "e2e-cascade-test"}, experiment_tags=["e2e", "cascade"], spans_per_trace=2, # root + 1 LLM child -> exercises parent_span_id remap feedback_scores_per_trace=[ {"name": "correctness", "value": 0.9, "reason": "matches reference"}, {"name": "latency_p95", "value": 230.5}, ], ) # Run the migration. audit_path = tmp_path / "audit.json" result = run_migrate_cli( [ "dataset", dataset_name, "--to-project", target_project_name, ], audit_log_path=str(audit_path), ) assert result.returncode == 0, result.stdout + result.stderr # Verify target: same version count, per-version content set-equal # under hash, display order matches at every version. target = rest.datasets.get_dataset_by_identifier( dataset_name=dataset_name, project_name=target_project_name ) tgt_versions = chronological_versions(rest, target.id) assert len(tgt_versions) == len(src_versions), ( f"target version count {len(tgt_versions)} != source {len(src_versions)} " "— Slice 2 contract requires N=N" ) for src_v, tgt_v, exp_hashes, exp_order in zip( src_versions, tgt_versions, expected_hashes, expected_orders ): items = stream_items_wire( rest, dataset_name=dataset_name, project_name=target_project_name, version_hash=tgt_v.version_hash, ) actual_hashes = item_hashes(items) actual_order = display_order(items) assert actual_hashes == exp_hashes, ( f"version {tgt_v.version_name}: target items don't match source. " f"Missing on target: {exp_hashes - actual_hashes}; " f"extra on target: {actual_hashes - exp_hashes}" ) assert actual_order == exp_order, ( f"version {tgt_v.version_name}: display order diverged " f"(source: {exp_order}, target: {actual_order})" ) # Audit log records one per-version entry per replayed source version. audit = json.loads(audit_path.read_text()) per_version_records = [ a for a in audit["actions"] if a["type"] == "replay_dataset_version" ] assert len(per_version_records) == len(src_versions) # Per-version deltas: v1=(3 adds), v2=(1 add, 1 mod), v3=(1 add, 1 mod, 1 del). assert ( per_version_records[0]["items_added"], per_version_records[0]["items_modified"], per_version_records[0]["items_deleted"], ) == (3, 0, 0) assert ( per_version_records[1]["items_added"], per_version_records[1]["items_modified"], per_version_records[1]["items_deleted"], ) == (1, 1, 0) assert ( per_version_records[2]["items_added"], per_version_records[2]["items_modified"], per_version_records[2]["items_deleted"], ) == (1, 1, 1) # ── Cascade fidelity ── # The destination project should now have a copy of the source # experiment with: a remapped dataset_version_id, fresh item ids # carrying remapped trace ids, traces+spans landing under the # destination project, feedback scores re-emitted on the destination # traces, and per-item write-side fidelity (input/output) preserved. dest_exp = find_destination_experiment( rest, destination_dataset_id=target.id, experiment_name=experiment_name, ) # FKs remapped. assert dest_exp.id != cascade_seed["experiment_id"] assert dest_exp.dataset_id == target.id # The destination experiment must reference one of the target # versions (the cascade picks the remap of v1). target_version_ids = {v.id for v in tgt_versions} assert dest_exp.dataset_version_id in target_version_ids # Items: one per source item, with FRESH trace ids (disjoint from # source). Per-item input/output/usage/cost are READ-ONLY on the BE # (computed/aggregated from the underlying trace + span entities); # we assert the trace + span fidelity below instead. dest_items = destination_experiment_items( rest, experiment_id=dest_exp.id, dataset_id=target.id, ) assert len(dest_items) == len(v1_item_ids) dest_trace_ids = {it.trace_id for it in dest_items} assert dest_trace_ids.isdisjoint(set(cascade_seed["trace_ids"])), ( "destination experiment items should reference new trace ids, " "not the source's" ) # Each destination trace exists under the target project and has the # same span shape as the source (root + 1 child = 2 spans). for new_trace_id in dest_trace_ids: dest_spans = destination_spans_for_trace( rest, trace_id=new_trace_id, project_name=target_project_name, ) assert len(dest_spans) == 2, ( f"trace {new_trace_id} should have 2 spans (root + child), " f"got {len(dest_spans)}" ) # Topological remap: exactly one root (parent_span_id=None), # the other span points at the root via parent_span_id. roots = [s for s in dest_spans if s.parent_span_id is None] assert len(roots) == 1, f"trace {new_trace_id} should have one root span" children = [s for s in dest_spans if s.parent_span_id is not None] assert all(c.parent_span_id == roots[0].id for c in children), ( f"trace {new_trace_id} child spans should remap parent_span_id " "to the new root id" ) # Trace-level feedback scores re-emitted on the destination trace. dest_scores = destination_feedback_scores_for_trace( rest, trace_id=new_trace_id ) score_names = {s.name for s in dest_scores} assert score_names == {"correctness", "latency_p95"}, ( f"trace {new_trace_id}: expected feedback score names " f"{{'correctness', 'latency_p95'}}, got {score_names}" ) # ── Deep-equal source vs. destination ── # Verify field-by-field that experiment + items + traces + spans # round-trip the cascade modulo remapped IDs. Pairing strategy: # both sides sorted by trace ``name`` (assigned by the seed as # "task-0", "task-1", "task-2" and carried verbatim through the # cascade), guaranteeing stable positional correspondence. src_exp = find_destination_experiment( rest, destination_dataset_id=source_id, experiment_name=experiment_name, ) src_items_compare = destination_experiment_items( rest, experiment_id=cascade_seed["experiment_id"], dataset_id=source_id, ) # Sort both sides by trace name for stable pairing. Build a # trace_id -> name map by reading each trace once. src_trace_names = { it.trace_id: rest.traces.get_trace_by_id(id=it.trace_id).name for it in src_items_compare } dst_trace_names = { it.trace_id: rest.traces.get_trace_by_id(id=it.trace_id).name for it in dest_items } src_items_compare.sort(key=lambda it: src_trace_names[it.trace_id]) dest_items_sorted = sorted( dest_items, key=lambda it: dst_trace_names[it.trace_id] ) src_trace_ids_sorted = [it.trace_id for it in src_items_compare] dst_trace_ids_sorted = [it.trace_id for it in dest_items_sorted] compare_cascade( rest_client=rest, source_experiment=src_exp, destination_experiment=dest_exp, source_item_ids=v1_item_ids, destination_item_ids=[it.dataset_item_id for it in dest_items_sorted], source_trace_ids=src_trace_ids_sorted, destination_trace_ids=dst_trace_ids_sorted, source_items_compare=src_items_compare, destination_items_compare=dest_items_sorted, ) class TestMigrateDatasetEnvironmentPreservation: """OPIK-6695: the cascade must preserve the ``environment`` column on traces, spans, and the BE-materialized ``trace_threads`` row. Pre-2026-05-07 rows default to ``''`` in ClickHouse and replay trivially; the regression this guards is the post-migration reset of a non-empty ``environment`` to ``''`` because the re-emit payload didn't carry the field. """ def test_environment_round_trips_on_traces_spans_and_threads( self, opik_client: opik.Opik, source_project_name: str, target_project_name: str, dataset_name: str, tmp_path: Path, ) -> None: from opik import id_helpers from opik.rest_api.types.dataset_item_write import DatasetItemWrite rest = opik_client.rest_client # Single-version dataset with two items -> two cascaded traces. source_id = create_dataset_shell(rest, dataset_name, source_project_name) rest.datasets.create_or_update_dataset_items( dataset_id=source_id, items=[ DatasetItemWrite(source="manual", data={"q": "Q1", "a": "A1"}), DatasetItemWrite(source="manual", data={"q": "Q2", "a": "A2"}), ], batch_group_id=id_helpers.generate_id(), ) v1 = rest.datasets.list_dataset_versions(id=source_id, page=1, size=1).content[ 0 ] v1_items = stream_items_wire( rest, dataset_name=dataset_name, project_name=source_project_name, version_hash=v1.version_hash, ) item_ids = [it.id for it in v1_items] # Seed: traces tagged environment="production" + grouped into one # thread; spans tagged environment="staging". The trace env and # span env differ deliberately so a single shared value couldn't # mask a per-entity bug, and the thread inherits the trace env. experiment_name = f"e2e-env-{random_chars()}" thread_id = f"env-thread-{random_chars()}" seed_experiment_with_trace_tree( rest, experiment_name=experiment_name, dataset_name=dataset_name, dataset_id=source_id, dataset_version_id=v1.id, project_name=source_project_name, item_ids=item_ids, spans_per_trace=2, trace_environment="production", span_environment="staging", thread_id=thread_id, ) audit_path = tmp_path / "audit.json" result = run_migrate_cli( ["dataset", dataset_name, "--to-project", target_project_name], audit_log_path=str(audit_path), ) assert result.returncode == 0, result.stdout + result.stderr target = rest.datasets.get_dataset_by_identifier( dataset_name=dataset_name, project_name=target_project_name ) dest_exp = find_destination_experiment( rest, destination_dataset_id=target.id, experiment_name=experiment_name, ) dest_items = destination_experiment_items( rest, experiment_id=dest_exp.id, dataset_id=target.id, ) assert len(dest_items) == len(item_ids) # (a) traces and (b) spans keep their source environment verbatim. for dest_item in dest_items: dest_trace = rest.traces.get_trace_by_id(id=dest_item.trace_id) assert dest_trace.environment == "production", ( f"trace {dest_item.trace_id} lost environment: " f"got {dest_trace.environment!r}" ) dest_spans = destination_spans_for_trace( rest, trace_id=dest_item.trace_id, project_name=target_project_name, ) assert dest_spans, f"trace {dest_item.trace_id} has no destination spans" assert all(span.environment == "staging" for span in dest_spans), ( "destination spans lost environment: " f"{[span.environment for span in dest_spans]}" ) # (c) the destination thread row -- materialized by the BE from the # cascaded traces -- inherits the same environment. Polls because # thread materialization is eventually consistent. assert synchronization.until( lambda: bool( opik_client.search_threads( project_name=target_project_name, filter_string=f'id = "{thread_id}"', ) ), max_try_seconds=30, ), f"destination thread {thread_id!r} never materialized" threads = opik_client.search_threads( project_name=target_project_name, filter_string=f'id = "{thread_id}"', ) assert len(threads) == 1 assert threads[0].environment == "production", ( f"destination thread {thread_id!r} lost environment: " f"got {threads[0].environment!r}" )