Add test on replicate nan

NREL · Oct 21, 2024 · 0946531 · 0946531
1 parent 166c939
commit 0946531
Showing 1 changed file with 36 additions and 0 deletions.
diff --git a/tests/model_tuning_test.py b/tests/model_tuning_test.py
@@ -28,6 +28,42 @@ def test_replicate_nan_values(self):
         assert result_df.equals(expected_df_2)
         assert df_1.equals(expected_df_1)
 
+    def test_replicate_nan_values_with_time(self):
+        # Sample dataframes
+        data_1 = {"A": [1, 2, np.nan, 4], "B": [5, np.nan, 7, 8], "C": [np.nan, 1, 1, 1]}
+        data_2 = {"A": [10, 20, 30, 40], "B": [50, 60, 70, 80]}
+        df_1 = pd.DataFrame(data_1)
+        df_2 = pd.DataFrame(data_2)
+
+        df_1["time"] = pd.date_range("2021-01-01", periods=4)
+        df_2["time"] = df_1["time"]
+
+        # Call the function to replicate NaN values
+        result_df = replicate_nan_values(df_1, df_2)
+
+        print(result_df)
+
+        # Expected output
+        expected_df_1 = pd.DataFrame(
+            {
+                "A": [1, 2, np.nan, 4],
+                "B": [5, np.nan, 7, 8],
+                "C": [np.nan, 1, 1, 1],
+                "time": pd.date_range("2021-01-01", periods=4),
+            }
+        )
+        expected_df_2 = pd.DataFrame(
+            {
+                "A": [10, 20, np.nan, 40],
+                "B": [50, np.nan, 70, 80],
+                "time": pd.date_range("2021-01-01", periods=4),
+            }
+        )
+
+        # Check if the result matches the expected output
+        assert result_df.equals(expected_df_2)
+        assert df_1.equals(expected_df_1)
+
     def test_evaluate_overall_wake_loss(self):
         # Create a sample DataFrame for testing
         data = {"pow_ref": [5, 2, 3], "pow_test": [1, 2, 2]}