mirror of
https://github.com/TauricResearch/TradingAgents.git
synced 2026-09-26 22:42:40 +03:00
Yijia-Xiao
b176c9374b
- dropping undated rows copies the frame before prices are coerced, so no chained-assignment write
82 lines
3.0 KiB
Python
82 lines
3.0 KiB
Python
"""Tests for tolerating a non-`Date` index column in the Yahoo OHLCV loader (#890).
|
|
|
|
Guards against a download frame whose date column is `index` or `Datetime`
|
|
instead of `Date`, which would otherwise silently drop every indicator.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import warnings
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
from tradingagents.dataflows.vendors.yahoo import ohlcv
|
|
|
|
|
|
def _ohlcv(date_col: str) -> pd.DataFrame:
|
|
"""OHLCV frame whose date column is named `date_col`."""
|
|
dates = pd.bdate_range("2026-04-01", periods=10)
|
|
return pd.DataFrame({
|
|
date_col: dates,
|
|
"Open": [100.0 + i for i in range(10)],
|
|
"High": [101.0 + i for i in range(10)],
|
|
"Low": [99.0 + i for i in range(10)],
|
|
"Close": [100.5 + i for i in range(10)],
|
|
"Volume": [1_000_000 + i for i in range(10)],
|
|
})
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestEnsureDateColumn:
|
|
def test_renames_index_column(self):
|
|
out = ohlcv._ensure_date_column(_ohlcv("index"))
|
|
assert "Date" in out.columns and "index" not in out.columns
|
|
|
|
def test_renames_datetime_and_date_variants(self):
|
|
assert "Date" in ohlcv._ensure_date_column(_ohlcv("Datetime")).columns
|
|
assert "Date" in ohlcv._ensure_date_column(_ohlcv("date")).columns
|
|
|
|
def test_leaves_existing_date_untouched(self):
|
|
df = _ohlcv("Date")
|
|
assert ohlcv._ensure_date_column(df) is df # no-op short-circuit
|
|
|
|
def test_no_datelike_column_is_left_alone(self):
|
|
df = pd.DataFrame({"Close": [1, 2, 3]})
|
|
out = ohlcv._ensure_date_column(df)
|
|
assert "Date" not in out.columns # nothing to rename; caller handles
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCleanDataframeAcrossVersions:
|
|
def test_clean_handles_index_column(self):
|
|
"""A frame with `index` instead of `Date` must still clean to a
|
|
usable, date-parsed frame (was KeyError: 'Date')."""
|
|
cleaned = ohlcv._clean_dataframe(_ohlcv("index"))
|
|
assert "Date" in cleaned.columns
|
|
assert pd.api.types.is_datetime64_any_dtype(cleaned["Date"])
|
|
assert len(cleaned) == 10
|
|
|
|
def test_clean_handles_legacy_date_column(self):
|
|
cleaned = ohlcv._clean_dataframe(_ohlcv("Date"))
|
|
assert len(cleaned) == 10
|
|
|
|
def test_indicators_compute_after_index_rename(self):
|
|
"""stockstats must compute indicators on a frame whose date column
|
|
arrived as `index`, instead of erroring per indicator."""
|
|
from stockstats import wrap
|
|
cleaned = ohlcv._clean_dataframe(_ohlcv("index"))
|
|
df = wrap(cleaned)
|
|
df["close_5_sma"] # triggers calculation
|
|
assert "close_5_sma" in df.columns
|
|
assert df["close_5_sma"].notna().any()
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_cleaning_a_frame_with_undated_rows_writes_to_its_own_copy():
|
|
raw = pd.DataFrame({"Date": ["2026-01-08", None, "2026-01-09"], "Close": ["1", "2", "x"]})
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("error")
|
|
cleaned = ohlcv._clean_dataframe(raw)
|
|
assert cleaned["Close"].tolist()[0] == 1.0
|