fix: failing tests

src/psycopt2d/load.py

@@ -338,9 +338,9 @@ def _drop_cols_and_rows_if_look_direction_not_met(
         for direction in ("ahead", "behind"):
 
             if direction in ("ahead", "behind"):
-                if self.cfg.data.min_lookahead_days:
+                if direction == "ahead":
                     n_days = self.cfg.data.min_lookahead_days
-                elif self.cfg.data.min_lookbehind_days:
+                elif direction == "behind":
                     n_days = self.cfg.data.min_lookbehind_days
                 else:
                     continue

src/psycopt2d/train_and_log_models.py

@@ -8,8 +8,6 @@
 - Run this script from project root with `python src/psycopt2d/train_and_log_models.py`
 """
 
-# TODO: Should be unified with the other train_and_log_models in application. Will be done when merging parent branch.
-
 import subprocess
 import time
 

tests/test_load.py

@@ -11,14 +11,13 @@ def test_load_lookbehind_exceeds_lookbehind_threshold():
     with initialize(version_base=None, config_path="../src/psycopt2d/config/"):
         cfg = compose(
             config_name="integration_testing.yaml",
+            overrides=["data.min_lookbehind_days=60"],
         )
 
         cfg = omegaconf_to_pydantic_objects(cfg)
-
-        cfg.data.min_lookahead_days = 30
         split_dataset = load_train_and_val_from_cfg(cfg)
 
-        assert split_dataset.train.shape[1] == 6
+        assert split_dataset.train.shape[1] == 7
 
 
 def test_load_lookbehind_not_in_lookbehind_combination():
@@ -27,11 +26,11 @@ def test_load_lookbehind_not_in_lookbehind_combination():
     with initialize(version_base=None, config_path="../src/psycopt2d/config/"):
         cfg = compose(
             config_name="integration_testing.yaml",
+            overrides=["data.lookbehind_combination=[30]"],
         )
 
         cfg = omegaconf_to_pydantic_objects(cfg)
 
-        cfg.data.lookbehind_combination = [30]
         split_dataset = load_train_and_val_from_cfg(cfg)
 
         assert split_dataset.train.shape[1] == 6

tests/test_train_model.py

@@ -3,6 +3,7 @@
 import pytest
 from hydra import compose, initialize
 
+from psycopt2d.load import load_train_from_cfg
 from psycopt2d.models import MODELS
 from psycopt2d.train_model import main
 from psycopt2d.utils.configs import omegaconf_to_pydantic_objects