gahow/quant
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: handle single-ticker yahoo panels

Browse Source
This commit is contained in:
Gahow Wang
2026-04-18 00:03:07 +08:00
parent 3abc51e3e3
commit 1edce83430
2 changed files with 29 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
tests/test_market_data.py
View File

@@ -118,6 +118,27 @@ class UpdateMarketDataTests(unittest.TestCase):
pd.testing.assert_frame_equal(panels["volume"], expected, check_freq=False)
pd.testing.assert_frame_equal(saved_volume, expected, check_freq=False)
def test_update_market_data_handles_single_ticker_multiindex_download(self):
dates = pd.to_datetime(["2024-01-02", "2024-01-03"])
raw = pd.DataFrame(
{
("Close", "AAA"): [10.0, 11.0],
("Volume", "AAA"): [1000.0, 1100.0],
},
index=dates,
)
raw.columns = pd.MultiIndex.from_tuples(raw.columns)
with tempfile.TemporaryDirectory() as tmpdir:
with mock.patch.object(data_manager, "DATA_DIR", tmpdir):
with mock.patch("data_manager.yf.download", return_value=raw):
panels = data_manager.update_market_data("us", ["AAA"], ["close", "volume"])
expected_close = pd.DataFrame({"AAA": [10.0, 11.0]}, index=dates)
expected_volume = pd.DataFrame({"AAA": [1000.0, 1100.0]}, index=dates)
pd.testing.assert_frame_equal(panels["close"], expected_close, check_freq=False)
pd.testing.assert_frame_equal(panels["volume"], expected_volume, check_freq=False)
if __name__ == "__main__":
unittest.main()