From d2baa40fb130b1bea76a482ebe5af49088ac13d3 Mon Sep 17 00:00:00 2001
From: Yuante Li <104308117+WinstonLiyt@users.noreply.github.com>
Date: Mon, 10 Feb 2025 20:18:01 +0800
Subject: [PATCH] add the none check for data loader & feature output (#577)

---
 .../coder/data_science/feature/eval_tests/feature_test.txt    | 4 ++++
 .../raw_data_loader/eval_tests/data_loader_test.txt           | 4 ++++
 2 files changed, 8 insertions(+)

diff --git a/rdagent/components/coder/data_science/feature/eval_tests/feature_test.txt b/rdagent/components/coder/data_science/feature/eval_tests/feature_test.txt
index 726c258a..a99d75ae 100644
--- a/rdagent/components/coder/data_science/feature/eval_tests/feature_test.txt
+++ b/rdagent/components/coder/data_science/feature/eval_tests/feature_test.txt
@@ -38,6 +38,10 @@ def get_column_list(data):
     return data.columns.tolist() if isinstance(data, pd.DataFrame) else None
 
 
+assert X is not None, "The feature engineering function returned None for X."
+assert y is not None, "The feature engineering function returned None for y."
+assert X_test is not None, "The feature engineering function returned None for X_test."
+
 assert get_length(X_test) == get_length(
     test_ids
 ), f"Mismatch in length of test images and test IDs: X_test ({get_length(X_test)}) and test_ids ({get_length(test_ids)})"
diff --git a/rdagent/components/coder/data_science/raw_data_loader/eval_tests/data_loader_test.txt b/rdagent/components/coder/data_science/raw_data_loader/eval_tests/data_loader_test.txt
index f9d4e757..750dfe11 100644
--- a/rdagent/components/coder/data_science/raw_data_loader/eval_tests/data_loader_test.txt
+++ b/rdagent/components/coder/data_science/raw_data_loader/eval_tests/data_loader_test.txt
@@ -26,6 +26,10 @@ def get_width(data):
 def get_column_list(data):
     return data.columns.tolist() if isinstance(data, pd.DataFrame) else None
 
+assert X is not None, "Training data (X) is None."
+assert y is not None, "Training labels (y) are None."
+assert X_test is not None, "Test data (X_test) is None."
+assert test_ids is not None, "Test IDs (test_ids) are None."
 
 assert get_length(X_test) == get_length(
     test_ids