In e2e tests, wait for ingested data to become available, and run a connectivity_test before e2e tests (#193)

* In e2e tests, wait for ingested data to become available, and run a connectivity_test before e2e tests

* Code review fix

* Move sleep :-)

* Add --capture=no to the connectivity test

* pre-commit

Author: gabrys

.github/workflows/tests-e2e.yml

@@ -60,7 +60,10 @@ jobs:
           NEPTUNE_QUERY_RETRY_SOFT_TIMEOUT: "120"
           TEST_MARKERS: ${{ matrix.test_markers }}
         run: |
-          pytest --retries=3 --retry-delay=2 --junitxml="test-results/test-e2e.xml" ${TEST_MARKERS:+-m "$TEST_MARKERS"} tests/e2e
+          pytest --capture=no tests/e2e/test_connectivity.py &&
+          pytest --retries=1 --retry-delay=2 --junitxml="test-results/test-e2e.xml" ${TEST_MARKERS:+-m "$TEST_MARKERS"} tests/e2e
+        # --capture=no is passed to the test_connectivity run, so that we can see logs from data ingestion
+        # This helps with debugging. At the very least, it shows the URL of the environment under test
 
       - name: Notify on failure
         if: failure() && github.event_name == 'schedule'

tests/e2e/data_ingestion.py

@@ -25,6 +25,9 @@
 import neptune_scale.types
 from neptune_api import AuthenticatedClient
 
+from neptune_query.internal.identifiers import ProjectIdentifier
+from neptune_query.internal.retrieval import search
+
 IngestionHistogram = neptune_scale.types.Histogram
 IngestionFile = neptune_scale.types.File
 
@@ -110,6 +113,30 @@ def get_run_by_run_id(self: IngestedProjectData, run_id: str) -> IngestedRunData
         raise ValueError(f"Run not found: {run_id}")
 
 
+def _wait_for_ingestion(
+    client: AuthenticatedClient, project_identifier: ProjectIdentifier, expected_data: ProjectData
+) -> None:
+    for attempt in range(20):
+        found_runs = 0
+
+        for page in search.fetch_run_sys_ids(
+            client=client,
+            project_identifier=project_identifier,
+            filter_=None,
+        ):
+            found_runs += len(page.items)
+
+        if found_runs == len(expected_data.runs):
+            return
+
+        # Next attempt in 2 seconds, please
+        sleep(2)
+
+    raise RuntimeError(
+        f"Timed out waiting for data ingestion, " f"found runs: {found_runs} out of expected: {len(expected_data.runs)}"
+    )
+
+
 def ingest_project(
     *,
     client: AuthenticatedClient,
@@ -143,6 +170,8 @@ def ingest_project(
                 project_data=project_data,
             )
 
+            _wait_for_ingestion(client=client, project_identifier=project_identifier, expected_data=project_data)
+
     return IngestedProjectData(
         project_identifier=project_identifier,
         ingested_runs=[
@@ -222,8 +251,6 @@ def _ingest_runs(runs_data: list[RunData], api_token: str, project_identifier: s
     for run in runs:
         run.close()
 
-    sleep(2)  # Extra wait to ensure data is available to query before proceeding
-
 
 def _get_all_steps(run_data: RunData) -> Iterable[float]:
     # Collect all unique steps

tests/e2e/test_connectivity.py

@@ -0,0 +1,28 @@
+import pytest
+
+from tests.e2e.conftest import EnsureProjectFunction
+from tests.e2e.data_ingestion import (
+    IngestedProjectData,
+    ProjectData,
+    RunData,
+)
+
+
+@pytest.fixture(scope="module")
+def project_1(ensure_project: EnsureProjectFunction) -> IngestedProjectData:
+    return ensure_project(
+        project_data=ProjectData(
+            runs=[
+                RunData(
+                    experiment_name="test_connectivity_experiment",
+                    run_id="test_connectivity_run",
+                    configs={"im_alive": 1},
+                ),
+            ]
+        )
+    )
+
+
+def test_connectivity(client, project_1: IngestedProjectData) -> None:
+    """A placeholder test to ensure connectivity to the Neptune server"""
+    assert True