From 65dc59a10e390f60917862cfd4239f306ad1059d Mon Sep 17 00:00:00 2001
From: Martin Bernstorff <martinbernstorff@gmail.com>
Date: Wed, 26 Oct 2022 12:20:10 +0200
Subject: [PATCH] fix: failing tests

---
 src/psycopt2d/load.py                 | 4 ++--
 src/psycopt2d/train_and_log_models.py | 2 --
 tests/test_load.py                    | 7 +++----
 tests/test_train_model.py             | 1 +
 4 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/src/psycopt2d/load.py b/src/psycopt2d/load.py
index 092b09ce..41958499 100644
--- a/src/psycopt2d/load.py
+++ b/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
diff --git a/src/psycopt2d/train_and_log_models.py b/src/psycopt2d/train_and_log_models.py
index 764b28af..3608733f 100644
--- a/src/psycopt2d/train_and_log_models.py
+++ b/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
 
diff --git a/tests/test_load.py b/tests/test_load.py
index 094fe3c6..90991383 100644
--- a/tests/test_load.py
+++ b/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
diff --git a/tests/test_train_model.py b/tests/test_train_model.py
index c6265627..21762d11 100644
--- a/tests/test_train_model.py
+++ b/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