diff --git a/lean/components/cloud/data_downloader.py b/lean/components/cloud/data_downloader.py index f4fe1040..0941b2ac 100644 --- a/lean/components/cloud/data_downloader.py +++ b/lean/components/cloud/data_downloader.py @@ -11,6 +11,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +import shutil +import tempfile from pathlib import Path from datetime import datetime, timedelta from typing import Any, List, Callable @@ -23,8 +26,30 @@ def _store_local_file(file_content: bytes, file_path: Path): file_path.parent.mkdir(parents=True, exist_ok=True) - with file_path.open("wb+") as f: - f.write(file_content) + # Write to a temporary file first so a failed write never truncates the existing file + fd, temp_name = tempfile.mkstemp(dir=file_path.parent, prefix=f".{file_path.name}.", suffix=".tmp") + try: + with os.fdopen(fd, "wb") as f: + f.write(file_content) + if file_path.exists(): + shutil.copymode(file_path, temp_name) + os.replace(temp_name, file_path) + except BaseException: + if os.path.exists(temp_name): + os.remove(temp_name) + raise + + +def _validate_market_hours_database(file_content: bytes) -> None: + """Raises if the content is not a market hours database that MarketHoursDatabase can read.""" + from re import sub, DOTALL + from json import loads + + # Same parsing as MarketHoursDatabase._get_all_entries() + text = sub(r"(/\*.*\*/)", "", file_content.decode("utf-8"), flags=DOTALL) + database = loads(text) + if not isinstance(database, dict) or not isinstance(database.get("entries"), dict): + raise ValueError("missing 'entries' object") def parse_timedelta(database_update_frequency: str): @@ -87,9 +112,15 @@ def update_database_files(self): _store_local_file(self._api_client.data.download_public_file( "https://raw.githubusercontent.com/QuantConnect/Lean/master/Data/symbol-properties/symbol-properties-database.csv"), data_dir / "symbol-properties" / "symbol-properties-database.csv") - _store_local_file(self._api_client.data.download_public_file( - "https://raw.githubusercontent.com/QuantConnect/Lean/master/Data/market-hours/market-hours-database.json"), - data_dir / "market-hours" / "market-hours-database.json") + market_hours_database = self._api_client.data.download_public_file( + "https://raw.githubusercontent.com/QuantConnect/Lean/master/Data/market-hours/market-hours-database.json") + try: + _validate_market_hours_database(market_hours_database) + except ValueError as e: + # Keep the existing local file; not a ValueError for the config option handler below + raise RuntimeError(f"Downloaded market-hours database is invalid ({e}), " + f"keeping the existing local file") from e + _store_local_file(market_hours_database, data_dir / "market-hours" / "market-hours-database.json") except MoreInfoError as e: if "not found" in str(e): pass diff --git a/tests/components/cloud/test_data_downloader.py b/tests/components/cloud/test_data_downloader.py new file mode 100644 index 00000000..282db79d --- /dev/null +++ b/tests/components/cloud/test_data_downloader.py @@ -0,0 +1,68 @@ +from datetime import datetime, timedelta +from pathlib import Path +from unittest import mock + +import pytest + +from lean.components.cloud.data_downloader import DataDownloader, _store_local_file + +VALID_DATABASE = '{"entries": {}}' + + +@pytest.fixture +def tmp_path() -> Path: + # The test suite runs on a fake filesystem (see conftest.py), so use a fixed path instead of pytest's tmp_path + path = Path("/data") + path.mkdir(parents=True) + return path + + +def _create_downloader(data_dir: Path, market_hours: bytes) -> DataDownloader: + lean_config_manager = mock.Mock() + lean_config_manager.get_lean_config.return_value = {} + lean_config_manager.get_data_directory.return_value = data_dir + api_client = mock.Mock() + api_client.data.download_public_file.side_effect = \ + lambda url: market_hours if url.endswith(".json") else b"symbol,properties\n" + return DataDownloader(mock.Mock(), api_client, lean_config_manager, "0.00:00:01") + + +def _market_hours_file(data_dir: Path) -> Path: + path = data_dir / "market-hours" / "market-hours-database.json" + path.parent.mkdir(parents=True) + path.write_text(VALID_DATABASE, encoding="utf-8") + return path + + +def test_update_database_files_keeps_existing_file_when_response_is_malformed(tmp_path: Path) -> None: + path = _market_hours_file(tmp_path) + + _create_downloader(tmp_path, b"not json").update_database_files() + + assert path.read_text(encoding="utf-8") == VALID_DATABASE + + +def test_update_database_files_keeps_existing_file_when_entries_are_missing(tmp_path: Path) -> None: + path = _market_hours_file(tmp_path) + + _create_downloader(tmp_path, b'{"other": 1}').update_database_files() + + assert path.read_text(encoding="utf-8") == VALID_DATABASE + + +def test_update_database_files_replaces_file_with_valid_commented_response(tmp_path: Path) -> None: + path = _market_hours_file(tmp_path) + new_content = b'/* comment */ {"entries": {"Equity-usa-[*]": {}}}' + + _create_downloader(tmp_path, new_content).update_database_files() + + assert path.read_bytes() == new_content + + +def test_store_local_file_leaves_no_temporary_files(tmp_path: Path) -> None: + path = tmp_path / "dir" / "file.json" + + _store_local_file(b"content", path) + + assert path.read_bytes() == b"content" + assert [p.name for p in path.parent.iterdir()] == ["file.json"]