Make ensure_project itself a fixture (#134)

* Make ensure_project itself a fixture

* Update e2e/experimental/test_fetch_experiments_table as well

* Add docstring to ensure_project

* Update whitespace in docstring

* code review notes - doc update

Co-authored-by: Piotr Łusakowski <[email protected]>

* Further docs update for ensure_project fixture

---------

Co-authored-by: Piotr Łusakowski <[email protected]>

Author: gabrys

tests/e2e/conftest.py

@@ -37,6 +37,11 @@
     PATH,
     TEST_DATA,
 )
+from tests.e2e.data_ingestion import (
+    IngestedProjectData,
+    ProjectData,
+    ingest_project,
+)
 
 
 @dataclass
@@ -236,3 +241,60 @@ def temp_dir():
 
 def extract_pages(generator):
     return list(it.chain.from_iterable(i.items for i in generator))
+
+
+class EnsureProjectFunction:
+    """
+    See ensure_project fixture docstring for usage details.
+    """
+
+    def __init__(self, client: AuthenticatedClient, api_token: str, workspace: str, test_execution_id: str):
+        self.client = client
+        self.api_token = api_token
+        self.workspace = workspace
+        self.test_execution_id = test_execution_id
+
+    def __call__(self, project_data: ProjectData, unique_key: str | None = None) -> IngestedProjectData:
+        return ingest_project(
+            client=self.client,
+            api_token=self.api_token,
+            workspace=self.workspace,
+            unique_key=unique_key or self.test_execution_id,
+            project_data=project_data,
+        )
+
+
+@pytest.fixture(scope="session")
+def ensure_project(client, api_token, workspace, test_execution_id) -> EnsureProjectFunction:
+    """Fixture returning a function-like object that can be used to create or retrieve projects with specified data.
+
+    Arguments for the returned callable:
+        project_data: Data to initialize the project with
+        unique_key: [optional] A unique identifier for tagging and later filtering projects, experiments, and runs.
+                    Tests spanning multiple projects should provide a unique_key that embeds the test_execution_id
+                    plus any additional context needed for disambiguation (e.g., module name).
+                    If not provided, the unique test_execution_id will be used.
+
+    Returns:
+        IngestedProjectData containing information about the created/retrieved project
+
+    Example:
+        @pytest.fixture
+        def project_gamma(ensure_project):
+            ingested_project = ensure_project(
+                ProjectData(
+                    project_name_base="test_project",
+                    runs=[
+                        RunData(
+                            experiment_name_base="experiment_gamma",
+                            config={"param": "value"},
+                            float_series={"metrics/accuracy": {0: 0.5, 1: 0.75, 2: 0.9}},
+                            string_series={"logs/status": {0: "started", 1: "in_progress", 2: "completed"}},
+                            histogram_series={ ... },
+                            files={ ... },
+                        )
+                        ...
+                    ],
+                ))
+    """
+    return EnsureProjectFunction(client, api_token, workspace, test_execution_id)

tests/e2e/data_ingestion.py

@@ -128,7 +128,7 @@ class IngestedProjectData:
     ingested_runs: list[IngestedRunData]
 
 
-def ensure_project(
+def ingest_project(
     *,
     client: AuthenticatedClient,
     api_token: str,

tests/e2e/experimental/test_fetch_experiments_table_global.py

@@ -16,7 +16,6 @@
     IngestedRunData,
     ProjectData,
     RunData,
-    ensure_project,
 )
 
 pytestmark = pytest.mark.filterwarnings("ignore:.*fetch_.*_table.*:neptune_query.warnings.ExperimentalWarning")
@@ -248,7 +247,7 @@ def test_fetch_experiments_table_with_empty_attributes(project_1, project_2, uni
 
 
 @pytest.fixture(scope="session")
-def project_1(client, api_token, workspace, unique_execution_module_key) -> IngestedProjectData:
+def project_1(ensure_project, unique_execution_module_key) -> IngestedProjectData:
     exp_project_1_root = RunData(
         experiment_name_base="exp_project_1",
         run_id_base="run_project_1_root",
@@ -317,9 +316,6 @@ def project_1(client, api_token, workspace, unique_execution_module_key) -> Inge
     )
 
     return ensure_project(
-        client=client,
-        api_token=api_token,
-        workspace=workspace,
         unique_key=unique_execution_module_key,
         project_data=ProjectData(
             project_name_base="global_fetch_experiments_table_project_1",
@@ -335,7 +331,7 @@ def project_1(client, api_token, workspace, unique_execution_module_key) -> Inge
 
 
 @pytest.fixture(scope="session")
-def project_2(client, api_token, workspace, unique_execution_module_key) -> IngestedProjectData:
+def project_2(ensure_project, unique_execution_module_key) -> IngestedProjectData:
     exp_project_2_root = RunData(
         experiment_name_base="exp_project_2",
         run_id_base="run_project_2_root",
@@ -404,9 +400,6 @@ def project_2(client, api_token, workspace, unique_execution_module_key) -> Inge
     )
 
     return ensure_project(
-        client=client,
-        api_token=api_token,
-        workspace=workspace,
         unique_key=unique_execution_module_key,
         project_data=ProjectData(
             project_name_base="global_fetch_experiments_table_project_2",

tests/e2e/experimental/test_fetch_runs_table_global.py

@@ -16,7 +16,6 @@
     IngestedRunData,
     ProjectData,
     RunData,
-    ensure_project,
 )
 
 pytestmark = pytest.mark.filterwarnings("ignore:.*fetch_.*_table.*:neptune_query.warnings.ExperimentalWarning")
@@ -243,11 +242,8 @@ def test_fetch_runs_table_with_empty_attributes(project_1, project_2, unique_exe
 
 
 @pytest.fixture(scope="session")
-def project_1(client, api_token, workspace, unique_execution_module_key) -> IngestedProjectData:
+def project_1(ensure_project, unique_execution_module_key) -> IngestedProjectData:
     return ensure_project(
-        client=client,
-        api_token=api_token,
-        workspace=workspace,
         unique_key=unique_execution_module_key,
         project_data=ProjectData(
             project_name_base="project_1",
@@ -298,11 +294,8 @@ def project_1(client, api_token, workspace, unique_execution_module_key) -> Inge
 
 
 @pytest.fixture(scope="session")
-def project_2(client, api_token, workspace, unique_execution_module_key) -> IngestedProjectData:
+def project_2(ensure_project, unique_execution_module_key) -> IngestedProjectData:
     return ensure_project(
-        client=client,
-        api_token=api_token,
-        workspace=workspace,
         unique_key=unique_execution_module_key,
         project_data=ProjectData(
             project_name_base="project_2",

tests/e2e/internal/retrieval/test_series.py

@@ -27,7 +27,6 @@
     IngestionHistogram,
     ProjectData,
     RunData,
-    ensure_project,
     step_to_timestamp,
 )
 
@@ -75,13 +74,9 @@ def __eq__(self, other: NQInternalHistogram) -> bool:
 
 
 @pytest.fixture(scope="session")
-def project_1(client, api_token, workspace, test_execution_id) -> IngestedProjectData:
+def project_1(ensure_project) -> IngestedProjectData:
     return ensure_project(
-        client=client,
-        api_token=api_token,
-        workspace=workspace,
-        unique_key=test_execution_id,
-        project_data=ProjectData(
+        ProjectData(
             project_name_base="project_404",
             runs=[
                 RunData(