Rewrite internal/retrieval/test_metric_buckets.py

Author: pitercl

tests/e2e/internal/retrieval/test_metric_buckets.py

@@ -1,26 +1,88 @@
 import pytest
 
+from neptune_query.internal.filters import _Filter
 from neptune_query.internal.identifiers import (
     AttributeDefinition,
+    ProjectIdentifier,
     RunAttributeDefinition,
+    RunIdentifier,
+    SysId,
 )
+from neptune_query.internal.retrieval import search
 from neptune_query.internal.retrieval.metric_buckets import (
     TimeseriesBucket,
     fetch_time_series_buckets,
 )
 from neptune_query.internal.retrieval.search import ContainerType
-from tests.e2e.data import (
-    FLOAT_SERIES_PATHS,
-    PATH,
-    TEST_DATA,
+from tests.e2e.data_ingestion import (
+    IngestedProjectData,
+    ProjectData,
+    RunData,
 )
 from tests.e2e.metric_buckets import (
     aggregate_metric_buckets,
     calculate_global_range,
     calculate_metric_bucket_ranges,
 )
 
-EXPERIMENT = TEST_DATA.experiments[0]
+
+@pytest.fixture(scope="module")
+def project(ensure_project, test_execution_id) -> IngestedProjectData:
+    project_data = ProjectData(
+        project_name_base="metric-buckets-project",
+        runs=[
+            RunData(
+                experiment_name_base="metric-buckets-experiment",
+                run_id_base="metric-buckets-run-id",
+                configs={
+                    "configs/int-value": 7,
+                    "configs/string-value": "example-config",
+                },
+                float_series={
+                    "metrics/float-series-value_0": {
+                        0.0: 0.5,
+                        1.0: 1.5,
+                        2.0: 2.5,
+                        3.0: 3.5,
+                    },
+                    "metrics/float-series-value_1": {
+                        0.0: 10.0,
+                        1.0: 11.0,
+                        2.0: 12.0,
+                        3.0: 13.0,
+                    },
+                    "metrics/step": {
+                        0.0: 0.0,
+                        1.0: 1.0,
+                        2.0: 2.0,
+                        3.0: 3.0,
+                    },
+                },
+            )
+        ],
+    )
+
+    unique_key = f"{test_execution_id}__metric_buckets"
+    return ensure_project(project_data, unique_key=unique_key)
+
+
+@pytest.fixture(scope="module")
+def experiment_identifier(client, project) -> RunIdentifier:
+    project_identifier = ProjectIdentifier(project.project_identifier)
+    experiment_name = project.ingested_runs[0].experiment_name
+
+    sys_ids: list[SysId] = []
+    for page in search.fetch_experiment_sys_ids(
+        client=client,
+        project_identifier=project_identifier,
+        filter_=_Filter.name_eq(experiment_name),
+    ):
+        sys_ids.extend(page.items)
+
+    if len(sys_ids) != 1:
+        raise RuntimeError(f"Expected to fetch exactly one sys_id for {experiment_name}, got {sys_ids}")
+
+    return RunIdentifier(project_identifier=project_identifier, sys_id=SysId(sys_ids[0]))
 
 
 def test_fetch_time_series_buckets_does_not_exist(client, project, experiment_identifier):
@@ -43,15 +105,6 @@ def test_fetch_time_series_buckets_does_not_exist(client, project, experiment_id
     assert result == {run_definition: []}
 
 
-@pytest.mark.parametrize(
-    "attribute_name, expected_values",
-    [
-        (
-            FLOAT_SERIES_PATHS[0],
-            list(zip(EXPERIMENT.float_series[f"{PATH}/metrics/step"], EXPERIMENT.float_series[FLOAT_SERIES_PATHS[0]])),
-        ),
-    ],
-)
 @pytest.mark.parametrize(
     "limit",
     [2, 10, 100],
@@ -60,11 +113,15 @@ def test_fetch_time_series_buckets_does_not_exist(client, project, experiment_id
     "x_range",
     [None, (1, 2), (-100, 100)],
 )
-def test_fetch_time_series_buckets_single_series(
-    client, project, experiment_identifier, attribute_name, expected_values, limit, x_range
-):
+def test_fetch_time_series_buckets_single_series(client, project, experiment_identifier, limit, x_range):
     # given
-    run_definition = RunAttributeDefinition(experiment_identifier, AttributeDefinition(attribute_name, "float-series"))
+    run_definition = RunAttributeDefinition(
+        experiment_identifier, AttributeDefinition("metrics/float-series-value_0", "float-series")
+    )
+    ingested_run = project.ingested_runs[0]
+    step_values = ingested_run.float_series["metrics/step"]
+    series_values = ingested_run.float_series["metrics/float-series-value_0"]
+    expected_values = [(step_values[step], series_values[step]) for step in sorted(series_values.keys())]
 
     # when
     result = fetch_time_series_buckets(