turn on status code tests for string dps + fix bug (#1712)

cognitedata · Apr 10, 2024 · 2ace8a3 · 2ace8a3
1 parent 14a9731
commit 2ace8a3
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 15 deletions.
diff --git a/cognite/client/_api/datapoint_tasks.py b/cognite/client/_api/datapoint_tasks.py
@@ -622,7 +622,9 @@ def extract_nullable_raw_dps(dps: DatapointsRaw) -> list[float | str]:  # actual
         return list(map(DpsUnpackFns.nullable_raw_dp, dps))
 
     @staticmethod
-    def extract_nullable_raw_dps_numpy(dps: DatapointsRaw) -> tuple[npt.NDArray[np.float64], list[int]]:
+    def extract_nullable_raw_dps_numpy(
+        dps: DatapointsRaw, dtype: type[np.float64] | type[np.object_]
+    ) -> tuple[npt.NDArray[Any], list[int]]:
         # This is a very hot loop, thus we make some ugly optimizations:
         values = [None] * len(dps)
         missing: list[int] = []
@@ -632,7 +634,7 @@ def extract_nullable_raw_dps_numpy(dps: DatapointsRaw) -> tuple[npt.NDArray[np.f
             values[i] = dp  # type: ignore [call-overload]
             if dp is None:
                 add_missing(i)
-        arr = np.array(values, dtype=np.float64)
+        arr = np.array(values, dtype=dtype)
         return arr, missing
 
     @staticmethod
@@ -1206,15 +1208,15 @@ def _include_outside_points_in_result(self) -> None:
     def _unpack_and_store(self, idx: tuple[float, ...], dps: DatapointsRaw) -> None:  # type: ignore [override]
         if self.use_numpy:
             self.ts_data[idx].append(DpsUnpackFns.extract_timestamps_numpy(dps))
+            assert self.raw_dtype_numpy is not None
             if self.query.ignore_bad_datapoints:
-                assert self.raw_dtype_numpy is not None
                 self.dps_data[idx].append(DpsUnpackFns.extract_raw_dps_numpy(dps, self.raw_dtype_numpy))
             else:
                 # After this step, missing values (represented with None) will become NaNs and thus become
                 # indistinguishable from any NaNs that was returned! We need to store these timestamps in a property
                 # to allow our users to inspect them - but maybe even more important, allow the SDK to accurately
                 # use the DatapointsArray to replicate datapoints (exactly).
-                arr, missing_idxs = DpsUnpackFns.extract_nullable_raw_dps_numpy(dps)
+                arr, missing_idxs = DpsUnpackFns.extract_nullable_raw_dps_numpy(dps, self.raw_dtype_numpy)
                 self.dps_data[idx].append(arr)
                 if missing_idxs:
                     self.null_timestamps.update(self.ts_data[idx][-1][missing_idxs].tolist())

diff --git a/tests/tests_integration/test_api/test_datapoints.py b/tests/tests_integration/test_api/test_datapoints.py
@@ -698,8 +698,7 @@ def test_numpy_dtypes_conversions_for_string_and_numeric(self, cognite_client, a
         assert type(dp_dumped["timestamp"]) is int  # noqa: E721
         assert type(dp_dumped["value"]) is str  # noqa: E721
 
-    # @pytest.mark.parametrize("test_is_string", (True, False))
-    @pytest.mark.parametrize("test_is_string", (False,))  # TODO: Add string dps tests once feature hits greenfield
+    @pytest.mark.parametrize("test_is_string", (True, False))
     def test_n_dps_retrieved_with_without_uncertain_and_bad(self, retrieve_endpoints, ts_status_codes, test_is_string):
         if test_is_string:
             _, mixed_ts, _, bad_ts = ts_status_codes
@@ -729,8 +728,7 @@ def test_n_dps_retrieved_with_without_uncertain_and_bad(self, retrieve_endpoints
             assert len(bad2) == 0
             assert len(bad3) == 365
 
-    # @pytest.mark.parametrize("test_is_string", (True, False))
-    @pytest.mark.parametrize("test_is_string", (False,))  # TODO: Add string dps tests once feature hits greenfield
+    @pytest.mark.parametrize("test_is_string", (True, False))
     def test_outside_points_with_bad_and_uncertain(self, retrieve_endpoints, ts_status_codes, test_is_string):
         if test_is_string:
             _, mixed_ts, _, bad_ts = ts_status_codes
@@ -1334,13 +1332,8 @@ def test_retrieve_methods_in_target_unit(
             res = DatapointsArray(max=res.values)
         assert math.isclose(res.max[0], 212)
 
-    # @pytest.mark.parametrize("test_is_string", (True, False))
-    @pytest.mark.parametrize("test_is_string", (False,))  # TODO: Add string dps tests once feature hits greenfield
-    def test_status_codes_affect_aggregate_calculations(self, retrieve_endpoints, ts_status_codes, test_is_string):
-        if test_is_string:
-            _, mixed_ts, _, bad_ts = ts_status_codes
-        else:
-            mixed_ts, _, bad_ts, _ = ts_status_codes
+    def test_status_codes_affect_aggregate_calculations(self, retrieve_endpoints, ts_status_codes):
+        mixed_ts, _, bad_ts, _ = ts_status_codes  # No aggregates for string dps
         bad_xid = bad_ts.external_id
         for endpoint, uses_numpy in zip(retrieve_endpoints, (False, True)):
             dps_lst = endpoint(