chore: add force_data_generation environment variable (#69)

* chore: add force_data_generation environment variable

* fix: add force_data_generation environment variable to remaining tests

---------

Co-authored-by: Michał Sośnicki <michal.sosnicki@neptune.ai>

Author: michalsosn

.github/workflows/tests-e2e.yml

@@ -4,6 +4,12 @@ on:
   schedule:
     - cron: '0 8 * * *' # Run at 8:00 daily
   workflow_dispatch:
+    inputs:
+      force_data_generation:
+        description: 'Force data generation (skip experiment existence check)'
+        required: false
+        default: false
+        type: boolean
   push:
     branches:
       - main
@@ -51,6 +57,7 @@ jobs:
           NEPTUNE_E2E_WORKSPACE: ${{ secrets.E2E_WORKSPACE }}
           NEPTUNE_E2E_PROJECT: ${{ secrets[format('E2E_PROJECT{0}', matrix.env_suffix)] }}
           NEPTUNE_E2E_API_TOKEN: ${{ secrets[format('E2E_API_TOKEN{0}', matrix.env_suffix)] }}
+          NEPTUNE_E2E_FORCE_DATA_GENERATION: ${{ inputs.force_data_generation || 'false' }}
           TEST_MARKERS: ${{ matrix.test_markers }}
         run: |
           pytest --junitxml="test-results/test-e2e.xml" ${TEST_MARKERS:+-m "$TEST_MARKERS"} tests/e2e

tests/e2e/conftest.py

@@ -84,18 +84,21 @@ def project():
 @pytest.fixture(scope="module")
 def run_with_attributes(project, api_token, client):
     runs = {}
+    force_data_generation = os.getenv("NEPTUNE_E2E_FORCE_DATA_GENERATION", "").lower() in ("true", "1", "yes")
+
     for experiment in TEST_DATA.experiments:
         project_id = project.project_identifier
 
-        existing = next(
-            fetch_experiment_sys_attrs(
-                client,
-                identifiers.ProjectIdentifier(project_id),
-                _Filter.name_eq(experiment.name),
+        if not force_data_generation:
+            existing = next(
+                fetch_experiment_sys_attrs(
+                    client,
+                    identifiers.ProjectIdentifier(project_id),
+                    _Filter.name_eq(experiment.name),
+                )
             )
-        )
-        if existing.items:
-            continue
+            if existing.items:
+                continue
 
         run = Run(
             api_token=api_token,

tests/e2e/internal/composition/test_attributes.py

@@ -1,3 +1,4 @@
+import os
 import re
 import time
 from datetime import (
@@ -42,15 +43,17 @@ def run_with_attributes(client, api_token, project):
 
     project_identifier = project.project_identifier
 
-    existing = next(
-        fetch_experiment_sys_attrs(
-            client,
-            identifiers.ProjectIdentifier(project_identifier),
-            _Filter.name_eq(EXPERIMENT_NAME),
+    force_data_generation = os.getenv("NEPTUNE_E2E_FORCE_DATA_GENERATION", "").lower() in ("true", "1", "yes")
+    if not force_data_generation:
+        existing = next(
+            fetch_experiment_sys_attrs(
+                client,
+                identifiers.ProjectIdentifier(project_identifier),
+                _Filter.name_eq(EXPERIMENT_NAME),
+            )
         )
-    )
-    if existing.items:
-        return
+        if existing.items:
+            return
 
     run_id = str(uuid.uuid4())
 

tests/e2e/internal/composition/test_type_inference.py

@@ -1,3 +1,4 @@
+import os
 from datetime import (
     datetime,
     timezone,
@@ -35,15 +36,17 @@ def run_with_attributes(client, api_token, project):
 
     project_identifier = project.project_identifier
 
-    existing = next(
-        fetch_experiment_sys_attrs(
-            client,
-            identifiers.ProjectIdentifier(project_identifier),
-            _Filter.name_eq(EXPERIMENT_NAME),
+    force_data_generation = os.getenv("NEPTUNE_E2E_FORCE_DATA_GENERATION", "").lower() in ("true", "1", "yes")
+    if not force_data_generation:
+        existing = next(
+            fetch_experiment_sys_attrs(
+                client,
+                identifiers.ProjectIdentifier(project_identifier),
+                _Filter.name_eq(EXPERIMENT_NAME),
+            )
         )
-    )
-    if existing.items:
-        return
+        if existing.items:
+            return
 
     run_id = str(uuid.uuid4())