47 lines
2.0 KiB
Python
47 lines
2.0 KiB
Python
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
import pandas as pd
|
|
|
|
from research import fetch_historical
|
|
|
|
|
|
class FetchHistoricalTests(unittest.TestCase):
|
|
def test_fetch_all_historical_ohlcv_writes_field_specific_csvs(self):
|
|
dates = pd.to_datetime(["2024-01-02", "2024-01-03"])
|
|
raw = pd.DataFrame(
|
|
{
|
|
("Close", "AAA"): [10.0, 11.0],
|
|
("Close", "BBB"): [20.0, 21.0],
|
|
("High", "AAA"): [10.5, 11.5],
|
|
("High", "BBB"): [20.5, 21.5],
|
|
("Low", "AAA"): [9.5, 10.5],
|
|
("Low", "BBB"): [19.5, 20.5],
|
|
("Volume", "AAA"): [1000.0, 1100.0],
|
|
("Volume", "BBB"): [2000.0, 2100.0],
|
|
},
|
|
index=dates,
|
|
)
|
|
raw.columns = pd.MultiIndex.from_tuples(raw.columns)
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
with mock.patch.object(fetch_historical, "DATA_DIR", tmpdir):
|
|
with mock.patch.object(fetch_historical, "OUT_PATH", str(Path(tmpdir) / "us_pit.csv")):
|
|
with mock.patch("research.fetch_historical.uh.load_sp500_history", return_value={"AAA": [[None, None]], "BBB": [[None, None]]}):
|
|
with mock.patch("research.fetch_historical.uh.all_tickers_ever", return_value=["AAA", "BBB"]):
|
|
with mock.patch("research.fetch_historical.yf.download", return_value=raw):
|
|
panels = fetch_historical.fetch_all_historical_ohlcv(force=True)
|
|
|
|
self.assertEqual(set(panels.keys()), {"close", "high", "low", "volume"})
|
|
self.assertTrue((Path(tmpdir) / "us_pit.csv").exists())
|
|
self.assertTrue((Path(tmpdir) / "us_pit_close.csv").exists())
|
|
self.assertTrue((Path(tmpdir) / "us_pit_high.csv").exists())
|
|
self.assertTrue((Path(tmpdir) / "us_pit_low.csv").exists())
|
|
self.assertTrue((Path(tmpdir) / "us_pit_volume.csv").exists())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|