Merge pull request #211 from mobiusml/bugfix_changing_tmp_data_dir

Refactor Settings to Dynamically Set Temporary Directories

aana/configs/settings.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 
-from pydantic import BaseModel, field_validator
+from pydantic import BaseModel, model_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from aana.configs.db import DbSettings
@@ -55,10 +55,11 @@ class Settings(BaseSettings):
     """
 
     tmp_data_dir: Path = Path("/tmp/aana_data")  # noqa: S108
-    image_dir: Path = tmp_data_dir / "images"
-    video_dir: Path = tmp_data_dir / "videos"
-    audio_dir: Path = tmp_data_dir / "audios"
-    model_dir: Path = tmp_data_dir / "models"
+    image_dir: Path | None = None
+    video_dir: Path | None = None
+    audio_dir: Path | None = None
+    model_dir: Path | None = None
+
     num_workers: int = 2
 
     task_queue: TaskQueueSettings = TaskQueueSettings()
@@ -67,11 +68,24 @@ class Settings(BaseSettings):
 
     test: TestSettings = TestSettings()
 
-    @field_validator("tmp_data_dir", mode="after")
-    def create_tmp_data_dir(cls, path: Path) -> Path:
-        """Create the tmp_data_dir if it doesn't exist."""
-        path.mkdir(parents=True, exist_ok=True)
-        return path
+    @model_validator(mode="after")
+    def setup_resource_directories(self):
+        """Create the resource directories if they do not exist."""
+        if self.image_dir is None:
+            self.image_dir = self.tmp_data_dir / "images"
+        if self.video_dir is None:
+            self.video_dir = self.tmp_data_dir / "videos"
+        if self.audio_dir is None:
+            self.audio_dir = self.tmp_data_dir / "audios"
+        if self.model_dir is None:
+            self.model_dir = self.tmp_data_dir / "models"
+
+        self.tmp_data_dir.mkdir(parents=True, exist_ok=True)
+        self.image_dir.mkdir(parents=True, exist_ok=True)
+        self.video_dir.mkdir(parents=True, exist_ok=True)
+        self.audio_dir.mkdir(parents=True, exist_ok=True)
+        self.model_dir.mkdir(parents=True, exist_ok=True)
+        return self
 
     model_config = SettingsConfigDict(
         protected_namespaces=("settings", *pydantic_protected_fields),

aana/tests/units/test_settings.py

@@ -16,3 +16,39 @@ def test_custom_tmp_data_dir(monkeypatch):
     monkeypatch.setenv("TMP_DATA_DIR", test_path)
     settings = Settings()
     assert settings.tmp_data_dir == Path(test_path)
+
+
+def test_changing_tmp_data_dir():
+    """Test that changing the temporary data directory is reflected in the other directories."""
+    new_tmp_data_dir = Path("/tmp/new_tmp_data_dir")
+    settings = Settings(tmp_data_dir=new_tmp_data_dir)
+
+    assert settings.tmp_data_dir == new_tmp_data_dir
+    assert settings.image_dir == new_tmp_data_dir / "images"
+    assert settings.video_dir == new_tmp_data_dir / "videos"
+    assert settings.audio_dir == new_tmp_data_dir / "audios"
+    assert settings.model_dir == new_tmp_data_dir / "models"
+
+    # Check that we can change the image directory independently
+    new_image_dir = Path("/tmp/new_image_dir")
+    settings = Settings(tmp_data_dir=new_tmp_data_dir, image_dir=new_image_dir)
+    assert settings.tmp_data_dir == new_tmp_data_dir
+    assert settings.image_dir == new_image_dir
+
+    # Check that we can change the video directory independently
+    new_video_dir = Path("/tmp/new_video_dir")
+    settings = Settings(tmp_data_dir=new_tmp_data_dir, video_dir=new_video_dir)
+    assert settings.tmp_data_dir == new_tmp_data_dir
+    assert settings.video_dir == new_video_dir
+
+    # Check that we can change the audio directory independently
+    new_audio_dir = Path("/tmp/new_audio_dir")
+    settings = Settings(tmp_data_dir=new_tmp_data_dir, audio_dir=new_audio_dir)
+    assert settings.tmp_data_dir == new_tmp_data_dir
+    assert settings.audio_dir == new_audio_dir
+
+    # Check that we can change the model directory independently
+    new_model_dir = Path("/tmp/new_model_dir")
+    settings = Settings(tmp_data_dir=new_tmp_data_dir, model_dir=new_model_dir)
+    assert settings.tmp_data_dir == new_tmp_data_dir
+    assert settings.model_dir == new_model_dir