timmy / 股票資料下載與測試 (Stock Data Downloader & Testing)

import unittest from datetime import datetime, timedelta from unittest.mock import patch import pandas as pd import yfinance as yf from dateutil.relativedelta import relativedelta import error_printer error_printer.configure_pretty_errors() error_printer.configure_icecream() class StockDataDownloader: def __init__(self, months: int = 1): """ 初始化 StockDataDownloader 類的實例。 參數: months (int): 設定 start 日期為 end 日期之前的月數。預設值為 1。 """ self.months = months self.start_date, self.end_date = self.set_date_range() def set_date_range(self) -> tuple[datetime, datetime]: """ 設定日期範圍，用於下載股票資料。 返回: tuple[datetime, datetime]: 包含 start 和 end 日期的元組。 """ end = datetime.now() start = end - relativedelta(months=self.months) print("Start date:", start) print("End date:", end) return start, end def download_stock_data(self, stock_symbol: str) -> pd.DataFrame: """ 下載指定股票程式碼在指定日期範圍內的股票資料。 參數: stock_symbol (str): 要下載的股票程式碼。 返回: pd.DataFrame: 包含股票資料的 DataFrame。 """ return yf.download(stock_symbol, start=self.start_date, end=self.end_date) class TestStockDataDownloader(unittest.TestCase): def test_initialization(self): """測試類初始化是否正確設定月份和日期範圍。""" downloader = StockDataDownloader(3) self.assertEqual(downloader.months, 3) @patch(__name__ + ".datetime") def test_date_range(self, mock_datetime): """測試日期範圍是否正確計算。""" # 設定模擬的現在時間 mock_now = datetime(2024, 5, 6, 9, 27, 48, 439831) mock_datetime.now.return_value = mock_now downloader = StockDataDownloader(2) expected_start_date = mock_now - timedelta(days=60) # 假設每月30天 expected_end_date = mock_now actual_start_date, actual_end_date = downloader.set_date_range() self.assertEqual(actual_start_date, expected_start_date) self.assertEqual(actual_end_date, expected_end_date) @patch("yfinance.download") def test_download_stock_data(self, mock_download): """測試是否正確下載股票資料。""" # 設定模擬的現在時間 mock_now = datetime(2024, 5, 6, 9, 27, 48, 439831) mock_start_date = mock_now - timedelta(days=30) mock_end_date = mock_now # 設定模擬的股票資料 mock_data = pd.DataFrame({"Date": pd.date_range(start=mock_start_date, end=mock_end_date, freq="D"), "Open": [100.0] * 31, "High": [101.0] * 31, "Low": [99.0] * 31, "Close": [100.0] * 31, "Volume": [1000] * 31}) mock_data.set_index("Date", inplace=True) mock_download.return_value = mock_data downloader = StockDataDownloader(1) stock_data = downloader.download_stock_data("00929.TW") self.assertTrue(stock_data.equals(mock_data)) # 檢查目前程式是否被作為主程式執行 if __name__ == "__main__": unittest.main() # 建立 StockDataDownloader 實例 # downloader = StockDataDownloader(months=1) # 下載股票資料 # stock_data = downloader.download_stock_data("00929.TW") # 進一步處理股票資料 # ... # ic(stock_data)