diff --git a/tests/unit/test_stockstats_utils.py b/tests/unit/test_stockstats_utils.py new file mode 100644 index 00000000..0009bf07 --- /dev/null +++ b/tests/unit/test_stockstats_utils.py @@ -0,0 +1,37 @@ +import pandas as pd +from tradingagents.dataflows.stockstats_utils import _clean_dataframe + +def test_clean_dataframe_lowercases_columns(): + df = pd.DataFrame({ + "Date": ["2023-01-01", "2023-01-02", "2023-01-03"], + "Open": [10.0, 11.0, 12.0], + "HIGH": [10.5, 11.5, 12.5], + "low": [9.5, 10.5, 11.5], + "ClOsE": [10.2, 11.2, 12.2], + "Volume": [1000, 1100, 1200] + }) + + cleaned = _clean_dataframe(df) + + assert list(cleaned.columns) == ["date", "open", "high", "low", "close", "volume"] + assert len(cleaned) == 3 + +def test_clean_dataframe_handles_non_string_columns(): + df = pd.DataFrame({ + 1: [10.0, 11.0], + "Open": [10.0, 11.0] + }) + + cleaned = _clean_dataframe(df) + + assert list(cleaned.columns) == ["1", "open"] + +def test_clean_dataframe_does_not_mutate_original(): + df = pd.DataFrame({ + "Date": ["2023-01-01", "2023-01-02", "2023-01-03"], + "Open": [10.0, 11.0, 12.0] + }) + + _clean_dataframe(df) + + assert list(df.columns) == ["Date", "Open"] diff --git a/tradingagents/dataflows/stockstats_utils.py b/tradingagents/dataflows/stockstats_utils.py index 467156a2..ee0875b9 100644 --- a/tradingagents/dataflows/stockstats_utils.py +++ b/tradingagents/dataflows/stockstats_utils.py @@ -7,16 +7,10 @@ from .config import get_config def _clean_dataframe(data: pd.DataFrame) -> pd.DataFrame: - """Normalize a stock DataFrame for stockstats: parse dates, drop invalid rows, fill price gaps.""" - data["Date"] = pd.to_datetime(data["Date"], errors="coerce") - data = data.dropna(subset=["Date"]) - - price_cols = [c for c in ["Open", "High", "Low", "Close", "Volume"] if c in data.columns] - data[price_cols] = data[price_cols].apply(pd.to_numeric, errors="coerce") - data = data.dropna(subset=["Close"]) - data[price_cols] = data[price_cols].ffill().bfill() - - return data + """Ensure DataFrame has lowercase columns for stockstats.""" + df = data.copy() + df.columns = [str(c).lower() for c in df.columns] + return df class StockstatsUtils: