Test fetch_experiments_table

Author: pitercl

tests/e2e/v1/generator.py

@@ -248,9 +248,9 @@ def metrics_values(self, name: AttributeName) -> list[tuple[Step, Value]]:
         },
         metrics={
             "series-containing-inf": dict(zip(range(10), [inf, 1, -inf, 3, 4, inf, 6, -inf, 8, 9])),
-            "series-ending-with-inf": dict(zip(range(10), [1, 2, 3, 4, 5, 6, 7, 8, 9, inf])),
+            "series-ending-with-inf": dict(zip(range(10), [0, 1, 2, 3, 4, 5, 6, 7, 8, inf])),
             "series-containing-nan": dict(zip(range(10), [nan, 1, nan, 3, 4, nan, 6, nan, 8, 9])),
-            "series-ending-with-nan": dict(zip(range(10), [1, 2, 3, 4, 5, 6, 7, 8, 9, nan])),
+            "series-ending-with-nan": dict(zip(range(10), [0, 1, 2, 3, 4, 5, 6, 7, 8, nan])),
         },
     ),
 ]

tests/e2e/v1/test_experiments.py

@@ -19,6 +19,7 @@
     TEST_DATA,
     TEST_DATA_VERSION,
 )
+from tests.e2e.v1.generator import EXP_NAME_INF_NAN_RUN
 
 
 @pytest.mark.parametrize("sort_direction", ["asc", "desc"])
@@ -312,6 +313,38 @@ def test__fetch_experiments_table_with_attributes_regex_filter_for_metrics(
     assert df[expected.columns].columns.equals(expected.columns)
 
 
+@pytest.mark.skip(reason="Skipped until inf/nan handling is enabled in the backend")
+def test__fetch_experiments_table_nan_inf(new_project_id):
+    df = fetch_experiments_table(
+        project=new_project_id,
+        experiments=[EXP_NAME_INF_NAN_RUN],
+        attributes=[
+            "inf-float",
+            "nan-float",
+            "neg-inf-float",
+            "series-containing-inf",
+            "series-ending-with-inf",
+            "series-containing-nan",
+            "series-ending-with-nan",
+        ],
+    )
+
+    expected = pd.DataFrame(
+        {
+            "experiment": [EXP_NAME_INF_NAN_RUN],
+            "inf-float": [float("inf")],
+            "nan-float": [float("nan")],
+            "neg-inf-float": [float("-inf")],
+            "series-containing-inf": [9.0],
+            "series-ending-with-inf": [float("inf")],
+            "series-containing-nan": [9.0],
+            "series-ending-with-nan": [float("nan")],
+        }
+    ).set_index("experiment", drop=True)
+    assert df.shape == (1, 7)
+    pd.testing.assert_frame_equal(df[expected.columns], expected)
+
+
 @pytest.mark.parametrize(
     "arg_experiments, expected_subset",
     [

tests/e2e/v1/test_fetch_metrics.py

@@ -12,10 +12,6 @@
 import numpy as np
 import pandas as pd
 import pytest
-from e2e.v1.generator import (
-    RUN_BY_ID,
-    RUN_ID_INF_NAN_RUN,
-)
 
 from neptune_query import fetch_metrics
 from neptune_query.filters import (
@@ -41,6 +37,8 @@
 from tests.e2e.v1.generator import (
     EXP_NAME_INF_NAN_RUN,
     MULT_EXPERIMENT_HISTORY_EXP_2,
+    RUN_BY_ID,
+    RUN_ID_INF_NAN_RUN,
     timestamp_for_step,
 )