From fc0d7e618d67704b8ea9bca95e6877eebae3252d Mon Sep 17 00:00:00 2001 From: ciaranbor Date: Thu, 12 Mar 2026 17:57:19 +0000 Subject: [PATCH] Add storage config IO tests --- src/exo/shared/tests/test_storage.py | 116 +++++++++++++++++++++++++++ 1 file changed, 116 insertions(+) diff --git a/src/exo/shared/tests/test_storage.py b/src/exo/shared/tests/test_storage.py index ee6008c7e..0ae21d50f 100644 --- a/src/exo/shared/tests/test_storage.py +++ b/src/exo/shared/tests/test_storage.py @@ -1,4 +1,6 @@ from datetime import UTC, datetime +from pathlib import Path +from unittest.mock import patch from exo.shared.models.model_cards import ModelId from exo.shared.storage import ( @@ -8,6 +10,8 @@ from exo.shared.storage import ( decide_storage_action, get_download_rejected_events, get_lru_eviction_candidates, + load_storage_config, + persist_storage_config, ) from exo.shared.tests.conftest import get_pipeline_shard_metadata from exo.shared.types.common import NodeId @@ -516,3 +520,115 @@ class TestGetDownloadRejectedEvents: # Only InstanceDeleted, no TaskStatusUpdated for completed task assert len(events) == 1 assert isinstance(events[0], InstanceDeleted) + + +class TestPersistStorageConfig: + """Tests for persist_storage_config I/O.""" + + async def test_round_trip(self, tmp_path: Path) -> None: + cfg_file = tmp_path / "config.toml" + config = StorageConfig( + max_storage=Memory.from_gb(50), storage_policy="auto-evict" + ) + with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file): + await persist_storage_config(config) + loaded = await load_storage_config() + + assert loaded.storage_policy == "auto-evict" + assert loaded.max_storage is not None + assert abs(loaded.max_storage.in_gb - 50.0) < 0.1 + + async def test_preserves_other_keys(self, tmp_path: Path) -> None: + cfg_file = tmp_path / "config.toml" + cfg_file.write_text('some_other_key = "hello"\n') + config = StorageConfig(max_storage=Memory.from_gb(10)) + with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file): + await persist_storage_config(config) + + contents = cfg_file.read_text() + assert "some_other_key" in contents + assert "hello" in contents + assert "max_storage_gb" in contents + + async def test_clears_max_storage_gb_when_unlimited(self, tmp_path: Path) -> None: + cfg_file = tmp_path / "config.toml" + cfg_file.write_text('max_storage_gb = 50\nstorage_policy = "auto-evict"\n') + config = StorageConfig(max_storage=None, storage_policy="manual") + with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file): + await persist_storage_config(config) + + contents = cfg_file.read_text() + assert "max_storage_gb" not in contents + assert "manual" in contents + + async def test_creates_file_if_missing(self, tmp_path: Path) -> None: + cfg_file = tmp_path / "subdir" / "config.toml" + config = StorageConfig(max_storage=Memory.from_gb(25)) + with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file): + await persist_storage_config(config) + + assert cfg_file.exists() + assert "max_storage_gb" in cfg_file.read_text() + + +class TestLoadStorageConfig: + """Tests for load_storage_config I/O.""" + + async def test_defaults_when_empty_file(self, tmp_path: Path) -> None: + cfg_file = tmp_path / "config.toml" + cfg_file.write_text("") + with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file): + config = await load_storage_config() + + assert config.max_storage is None + assert config.storage_policy == "manual" + + async def test_reads_from_file(self, tmp_path: Path) -> None: + cfg_file = tmp_path / "config.toml" + cfg_file.write_text('max_storage_gb = 30.0\nstorage_policy = "auto-evict"\n') + with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file): + config = await load_storage_config() + + assert config.storage_policy == "auto-evict" + assert config.max_storage is not None + assert abs(config.max_storage.in_gb - 30.0) < 0.1 + + async def test_cli_overrides_file(self, tmp_path: Path) -> None: + cfg_file = tmp_path / "config.toml" + cfg_file.write_text('max_storage_gb = 30.0\nstorage_policy = "manual"\n') + with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file): + config = await load_storage_config( + max_storage_gb=100.0, storage_policy="auto-evict" + ) + + assert config.storage_policy == "auto-evict" + assert config.max_storage is not None + assert abs(config.max_storage.in_gb - 100.0) < 0.1 + + async def test_partial_cli_override(self, tmp_path: Path) -> None: + """CLI overrides only the fields provided, file values used for the rest.""" + cfg_file = tmp_path / "config.toml" + cfg_file.write_text('max_storage_gb = 30.0\nstorage_policy = "auto-evict"\n') + with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file): + config = await load_storage_config(max_storage_gb=50.0) + + assert config.storage_policy == "auto-evict" # from file + assert config.max_storage is not None + assert abs(config.max_storage.in_gb - 50.0) < 0.1 # from CLI + + async def test_defaults_on_corrupt_file(self, tmp_path: Path) -> None: + cfg_file = tmp_path / "config.toml" + cfg_file.write_text("not valid toml {{{") + with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file): + config = await load_storage_config() + + assert config.max_storage is None + assert config.storage_policy == "manual" + + async def test_creates_file_if_missing(self, tmp_path: Path) -> None: + cfg_file = tmp_path / "subdir" / "config.toml" + with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file): + config = await load_storage_config() + + assert config.max_storage is None + assert cfg_file.exists()