From 859cf7dac5b06bb7bf86fc139d3bc9ba1c088acf Mon Sep 17 00:00:00 2001 From: Ariel Ebersberger Date: Mon, 10 Jun 2024 13:55:13 +0200 Subject: [PATCH] Refactor Tests --- .../pipeline/model/training/torch_trainer.py | 10 +- tests/_core/_caching/test__cacher.py | 90 +++++------- tests/_core/_logging/test__logger.py | 1 + tests/conftest.py | 15 ++ tests/constants.py | 3 + tests/pipeline/model/test_model.py | 50 ++++--- .../augmentation/test_image_augmentations.py | 3 +- .../test_time_series_augmentations.py | 4 +- .../model/training/augmentation/test_utils.py | 3 +- .../model/training/models/test_timm.py | 3 +- .../model/training/test_pretrain_block.py | 1 + .../model/training/test_torch_trainer.py | 129 +++++++++--------- .../pipeline/model/training/test_training.py | 74 +++++----- .../model/training/test_training_block.py | 17 ++- .../transformation/test_transformation.py | 63 +++++---- .../test_transformation_block.py | 20 +-- tests/pipeline/test_ensemble.py | 45 +++--- tests/util.py | 18 --- 18 files changed, 280 insertions(+), 269 deletions(-) create mode 100644 tests/conftest.py create mode 100644 tests/constants.py delete mode 100644 tests/util.py diff --git a/epochalyst/pipeline/model/training/torch_trainer.py b/epochalyst/pipeline/model/training/torch_trainer.py index cad6651..9a259e7 100644 --- a/epochalyst/pipeline/model/training/torch_trainer.py +++ b/epochalyst/pipeline/model/training/torch_trainer.py @@ -484,12 +484,12 @@ def create_datasets( :return: The training and validation datasets. """ train_dataset = TensorDataset( - torch.tensor(x[train_indices]), - torch.tensor(y[train_indices]), + x[train_indices].clone().detach(), + y[train_indices].clone().detach(), ) test_dataset = TensorDataset( - torch.tensor(x[test_indices]), - torch.tensor(y[test_indices]), + x[test_indices].clone().detach(), + y[test_indices].clone().detach(), ) return train_dataset, test_dataset @@ -503,7 +503,7 @@ def create_prediction_dataset( :param x: The input data. :return: The prediction dataset. """ - return TensorDataset(torch.tensor(x)) + return TensorDataset(x.clone().detach()) def create_dataloaders( self, diff --git a/tests/_core/_caching/test__cacher.py b/tests/_core/_caching/test__cacher.py index d250ea7..3475087 100644 --- a/tests/_core/_caching/test__cacher.py +++ b/tests/_core/_caching/test__cacher.py @@ -1,15 +1,15 @@ import shutil -from epochalyst._core._caching._cacher import _Cacher -import numpy as np +from pathlib import Path + +import dask.array as da import dask.dataframe as dd +import numpy as np import pandas as pd import polars as pl -import dask.array as da -from tests.util import remove_cache_files import pytest -from pathlib import Path -TEMP_DIR = Path("tests/temp") +from epochalyst._core._caching._cacher import _Cacher +from tests.constants import TEMP_DIR class Implemented_Cacher(_Cacher): @@ -21,15 +21,8 @@ class Test_Cacher: cache_path = TEMP_DIR @pytest.fixture(autouse=True) - def run_around_tests(self): - # Code that will run before each test - TEMP_DIR.mkdir(exist_ok=True) - - yield - - # Code that will run after each - if TEMP_DIR.exists(): - shutil.rmtree(TEMP_DIR) + def run_always(self, setup_temp_dir): + pass def test_cacher_init(self): c = Implemented_Cacher() @@ -64,13 +57,13 @@ def test__cache_exists_storage_type_npy_exists(self): ) is True ) - remove_cache_files() def test__cache_exists_storage_type_parquet(self): c = Implemented_Cacher() assert ( c.cache_exists( - "test", {"storage_type": ".parquet", "storage_path": f"{self.cache_path}"} + "test", + {"storage_type": ".parquet", "storage_path": f"{self.cache_path}"}, ) is False ) @@ -81,11 +74,11 @@ def test__cache_exists_storage_type_parquet_exists(self): f.write("test") assert ( c.cache_exists( - "test", {"storage_type": ".parquet", "storage_path": f"{self.cache_path}"} + "test", + {"storage_type": ".parquet", "storage_path": f"{self.cache_path}"}, ) is True ) - remove_cache_files() def test__cache_exists_storage_type_csv(self): c = Implemented_Cacher() @@ -106,13 +99,13 @@ def test__cache_exists_storage_type_csv_exists(self): ) is True ) - remove_cache_files() def test__cache_exists_storage_type_npy_stack(self): c = Implemented_Cacher() assert ( c.cache_exists( - "test", {"storage_type": ".npy_stack", "storage_path": f"{self.cache_path}"} + "test", + {"storage_type": ".npy_stack", "storage_path": f"{self.cache_path}"}, ) is False ) @@ -123,11 +116,11 @@ def test__cache_exists_storage_type_npy_stack_exists(self): f.write("test") assert ( c.cache_exists( - "test", {"storage_type": ".npy_stack", "storage_path": f"{self.cache_path}"} + "test", + {"storage_type": ".npy_stack", "storage_path": f"{self.cache_path}"}, ) is True ) - remove_cache_files() def test__cache_exists_storage_type_pkl(self): c = Implemented_Cacher() @@ -148,13 +141,13 @@ def test__cache_exists_storage_type_pkl_exists(self): ) is True ) - remove_cache_files() def test__cache_exists_storage_type_unsupported(self): c = Implemented_Cacher() assert ( c.cache_exists( - "test", {"storage_type": ".new_type", "storage_path": f"{self.cache_path}"} + "test", + {"storage_type": ".new_type", "storage_path": f"{self.cache_path}"}, ) is False ) @@ -179,7 +172,9 @@ def test__store_cache_no_output_data_type(self): c = Implemented_Cacher() with pytest.raises(ValueError): c._store_cache( - "test", "test", {"storage_type": ".npy", "storage_path": f"{self.cache_path}"} + "test", + "test", + {"storage_type": ".npy", "storage_path": f"{self.cache_path}"}, ) # storage type .npy @@ -200,7 +195,6 @@ def test__store_cache_storage_type_npy_output_data_type_numpy_array(self): ) is True ) - remove_cache_files() def test__store_cache_storage_type_npy_output_data_type_dask_array(self): c = Implemented_Cacher() @@ -221,7 +215,6 @@ def test__store_cache_storage_type_npy_output_data_type_dask_array(self): ) is True ) - remove_cache_files() def test__store_cache_storage_type_npy_output_data_type_unsupported(self): c = Implemented_Cacher() @@ -252,11 +245,11 @@ def test__store_cache_storage_type_parquet_output_data_type_pandas_dataframe(sel ) assert ( c.cache_exists( - "test", {"storage_type": ".parquet", "storage_path": f"{self.cache_path}"} + "test", + {"storage_type": ".parquet", "storage_path": f"{self.cache_path}"}, ) is True ) - remove_cache_files() def test__store_cache_storage_type_parquet_output_data_type_dask_dataframe(self): c = Implemented_Cacher() @@ -274,11 +267,11 @@ def test__store_cache_storage_type_parquet_output_data_type_dask_dataframe(self) ) assert ( c.cache_exists( - "test", {"storage_type": ".parquet", "storage_path": f"{self.cache_path}"} + "test", + {"storage_type": ".parquet", "storage_path": f"{self.cache_path}"}, ) is True ) - remove_cache_files() def test__store_cache_storage_type_parquet_output_data_type_numpy_array(self): c = Implemented_Cacher() @@ -295,11 +288,11 @@ def test__store_cache_storage_type_parquet_output_data_type_numpy_array(self): ) assert ( c.cache_exists( - "test", {"storage_type": ".parquet", "storage_path": f"{self.cache_path}"} + "test", + {"storage_type": ".parquet", "storage_path": f"{self.cache_path}"}, ) is True ) - remove_cache_files() def test__store_cache_storage_type_parquet_output_data_type_dask_array(self): c = Implemented_Cacher() @@ -316,11 +309,11 @@ def test__store_cache_storage_type_parquet_output_data_type_dask_array(self): ) assert ( c.cache_exists( - "test", {"storage_type": ".parquet", "storage_path": f"{self.cache_path}"} + "test", + {"storage_type": ".parquet", "storage_path": f"{self.cache_path}"}, ) is True ) - remove_cache_files() def test__store_cache_storage_type_parquet_output_data_type_polars_dataframe(self): c = Implemented_Cacher() @@ -337,11 +330,11 @@ def test__store_cache_storage_type_parquet_output_data_type_polars_dataframe(sel ) assert ( c.cache_exists( - "test", {"storage_type": ".parquet", "storage_path": f"{self.cache_path}"} + "test", + {"storage_type": ".parquet", "storage_path": f"{self.cache_path}"}, ) is True ) - remove_cache_files() def test__store_cache_storage_type_parquet_output_data_type_unsupported(self): c = Implemented_Cacher() @@ -376,7 +369,6 @@ def test__store_cache_storage_type_csv_output_data_type_pandas_dataframe(self): ) is True ) - remove_cache_files() def test__store_cache_storage_type_csv_output_data_type_dask_dataframe(self): c = Implemented_Cacher() @@ -398,7 +390,6 @@ def test__store_cache_storage_type_csv_output_data_type_dask_dataframe(self): ) is True ) - remove_cache_files() def test__store_cache_storage_type_csv_output_data_type_polars_dataframe(self): c = Implemented_Cacher() @@ -419,7 +410,6 @@ def test__store_cache_storage_type_csv_output_data_type_polars_dataframe(self): ) is True ) - remove_cache_files() def test__store_cache_storage_type_csv_output_data_type_unsupported(self): c = Implemented_Cacher() @@ -450,7 +440,8 @@ def test__store_cache_storage_type_npy_stack_output_data_type_dask_array(self): ) assert ( c.cache_exists( - "test", {"storage_type": ".npy_stack", "storage_path": f"{self.cache_path}"} + "test", + {"storage_type": ".npy_stack", "storage_path": f"{self.cache_path}"}, ) is True ) @@ -501,7 +492,6 @@ def test__store_cache_storage_type_pkl_output_data_type_pandas_dataframe(self): ) is True ) - remove_cache_files() def test__store_cache_storage_type_pkl_output_data_type_dask_dataframe(self): c = Implemented_Cacher() @@ -523,7 +513,6 @@ def test__store_cache_storage_type_pkl_output_data_type_dask_dataframe(self): ) is True ) - remove_cache_files() # _get_cache def test__get_cache_no_cache_args(self): @@ -571,7 +560,6 @@ def test__get_cache_storage_type_npy_output_data_type_numpy_array(self): ) == "test" ) - remove_cache_files() def test__get_cache_storage_type_npy_output_data_type_dask_array(self): c = Implemented_Cacher() @@ -600,7 +588,6 @@ def test__get_cache_storage_type_npy_output_data_type_dask_array(self): .all() == x.compute().all() ) - remove_cache_files() def test__get_cache_storage_type_npy_output_data_type_unsupported(self): c = Implemented_Cacher() @@ -636,7 +623,6 @@ def test__get_cache_storage_type_parquet_output_data_type_pandas_dataframe(self) "output_data_type": "pandas_dataframe", }, ).equals(data) - remove_cache_files() def test__get_cache_storage_type_parquet_output_data_type_dask_dataframe(self): c = Implemented_Cacher() @@ -664,7 +650,6 @@ def test__get_cache_storage_type_parquet_output_data_type_dask_dataframe(self): .compute() .equals(data.compute()) ) - remove_cache_files() def test__get_cache_storage_type_parquet_output_data_type_numpy_array(self): c = Implemented_Cacher() @@ -688,7 +673,6 @@ def test__get_cache_storage_type_parquet_output_data_type_numpy_array(self): }, ) assert get_cache.all() == data.all() - remove_cache_files() def test__get_cache_storage_type_parquet_output_data_type_dask_array(self): c = Implemented_Cacher() @@ -712,7 +696,6 @@ def test__get_cache_storage_type_parquet_output_data_type_dask_array(self): }, ) assert get_cache.compute().all() == data.compute().all() - remove_cache_files() def test__get_cache_storage_type_parquet_output_data_type_polars_dataframe(self): c = Implemented_Cacher() @@ -736,7 +719,6 @@ def test__get_cache_storage_type_parquet_output_data_type_polars_dataframe(self) }, ) assert data.equals(get_cache) - remove_cache_files() def test__get_cache_storage_type_parquet_output_data_type_unsupported(self): c = Implemented_Cacher() @@ -773,7 +755,6 @@ def test__get_cache_storage_type_csv_output_data_type_pandas_dataframe(self): }, ) assert get_cache.equals(data) - remove_cache_files() def test__get_cache_storage_type_csv_output_data_type_dask_dataframe(self): c = Implemented_Cacher() @@ -798,7 +779,6 @@ def test__get_cache_storage_type_csv_output_data_type_dask_dataframe(self): }, ) assert get_cache.compute().reset_index(drop=True).equals(data.compute()) - remove_cache_files() def test__get_cache_storage_type_csv_output_data_type_polars_dataframe(self): c = Implemented_Cacher() @@ -822,7 +802,6 @@ def test__get_cache_storage_type_csv_output_data_type_polars_dataframe(self): }, ) assert data.equals(get_cache) - remove_cache_files() def test__get_cache_storage_type_csv_output_data_type_unsupported(self): c = Implemented_Cacher() @@ -859,7 +838,6 @@ def test__get_cache_storage_type_npy_stack_output_data_type_dask_array(self): }, ) assert get_cache.compute().all() == data.compute().all() - remove_cache_files() def test__get_cache_storage_type_npy_stack_output_data_type_unsupported(self): c = Implemented_Cacher() @@ -908,7 +886,6 @@ def test__get_cache_storage_type_pkl_output_data_type_pandas_dataframe(self): }, ) assert get_cache.equals(data) - remove_cache_files() def test__get_cache_storage_type_pkl_output_data_type_dask_dataframe(self): c = Implemented_Cacher() @@ -933,4 +910,3 @@ def test__get_cache_storage_type_pkl_output_data_type_dask_dataframe(self): }, ) assert get_cache.compute().reset_index(drop=True).equals(data.compute()) - remove_cache_files() diff --git a/tests/_core/_logging/test__logger.py b/tests/_core/_logging/test__logger.py index f422de9..1b4960f 100644 --- a/tests/_core/_logging/test__logger.py +++ b/tests/_core/_logging/test__logger.py @@ -1,4 +1,5 @@ import pytest + from epochalyst._core._logging._logger import _Logger diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..63f8ee3 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,15 @@ +import shutil + +import pytest + +from tests.constants import TEMP_DIR + + +@pytest.fixture +def setup_temp_dir(): + TEMP_DIR.mkdir(exist_ok=True) + + yield + + if TEMP_DIR.exists(): + shutil.rmtree(TEMP_DIR) diff --git a/tests/constants.py b/tests/constants.py new file mode 100644 index 0000000..7813510 --- /dev/null +++ b/tests/constants.py @@ -0,0 +1,3 @@ +from pathlib import Path + +TEMP_DIR = Path("tests/temp") diff --git a/tests/pipeline/model/test_model.py b/tests/pipeline/model/test_model.py index 0451f82..6c46b47 100644 --- a/tests/pipeline/model/test_model.py +++ b/tests/pipeline/model/test_model.py @@ -1,16 +1,19 @@ -from epochalyst.pipeline.model.model import ModelPipeline -import numpy as np -from tests.util import remove_cache_files +import shutil +from pathlib import Path from typing import Any -from epochalyst.pipeline.model.transformation.transformation import ( - TransformationPipeline, -) -from epochalyst.pipeline.model.transformation.transformation_block import ( - TransformationBlock, -) +import numpy as np +import pytest + +from epochalyst.pipeline.model.model import ModelPipeline +from epochalyst.pipeline.model.transformation.transformation import \ + TransformationPipeline +from epochalyst.pipeline.model.transformation.transformation_block import \ + TransformationBlock +from tests.constants import TEMP_DIR -class TestTransformation(TransformationBlock): + +class ExampleTransformation(TransformationBlock): def log_to_debug(self, message: str) -> None: pass @@ -18,7 +21,7 @@ def custom_transform(self, data: Any, **transform_args: Any) -> Any: return data * 2 -class TestTransformationPipeline(TransformationPipeline): +class ExampleTransformationPipeline(TransformationPipeline): def log_to_debug(self, message: str) -> None: pass @@ -27,6 +30,12 @@ def log_to_terminal(self, message: str) -> None: class TestModel: + cache_path = TEMP_DIR + + @pytest.fixture(autouse=True) + def run_always(self, setup_temp_dir): + pass + def test_model_pipeline_init(self): model = ModelPipeline() assert model is not None @@ -42,19 +51,18 @@ def test_model_pipeline_predict(self): assert model.predict(None) is None def test_model_get_x_y_cache_exists(self): - test_transformation = TestTransformation() - x_sys = TestTransformationPipeline([test_transformation]) + test_transformation = ExampleTransformation() + x_sys = ExampleTransformationPipeline([test_transformation]) model = ModelPipeline(x_sys=x_sys) cache_args = { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } assert model.get_x_cache_exists(cache_args) is False assert model.get_y_cache_exists(cache_args) is False - remove_cache_files() def test_model_get_x_y_cache_systems_none(self): model = ModelPipeline() @@ -62,22 +70,21 @@ def test_model_get_x_y_cache_systems_none(self): cache_args = { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } assert model.get_x_cache_exists(cache_args) is False assert model.get_y_cache_exists(cache_args) is False - remove_cache_files() def test_model_get_x_y_cache_exists_true(self): - test_transformation = TestTransformation() - x_sys = TestTransformationPipeline([test_transformation]) - y_sys = TestTransformationPipeline([test_transformation]) + test_transformation = ExampleTransformation() + x_sys = ExampleTransformationPipeline([test_transformation]) + y_sys = ExampleTransformationPipeline([test_transformation]) model = ModelPipeline(x_sys=x_sys, y_sys=y_sys) cache_args = { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } train_args = { @@ -94,4 +101,3 @@ def test_model_get_x_y_cache_exists_true(self): assert model.get_x_cache_exists(cache_args) is True assert model.get_y_cache_exists(cache_args) is True - remove_cache_files() diff --git a/tests/pipeline/model/training/augmentation/test_image_augmentations.py b/tests/pipeline/model/training/augmentation/test_image_augmentations.py index fb02bbb..323b880 100644 --- a/tests/pipeline/model/training/augmentation/test_image_augmentations.py +++ b/tests/pipeline/model/training/augmentation/test_image_augmentations.py @@ -1,6 +1,7 @@ -from epochalyst.pipeline.model.training.augmentation import image_augmentations import torch +from epochalyst.pipeline.model.training.augmentation import image_augmentations + class TestImageAugmentations: def test_cutmix(self): diff --git a/tests/pipeline/model/training/augmentation/test_time_series_augmentations.py b/tests/pipeline/model/training/augmentation/test_time_series_augmentations.py index 0c38494..122d2ec 100644 --- a/tests/pipeline/model/training/augmentation/test_time_series_augmentations.py +++ b/tests/pipeline/model/training/augmentation/test_time_series_augmentations.py @@ -1,7 +1,9 @@ import numpy as np -from epochalyst.pipeline.model.training.augmentation import time_series_augmentations import torch +from epochalyst.pipeline.model.training.augmentation import \ + time_series_augmentations + def set_torch_seed(seed: int = 42) -> None: """Set torch seed for reproducibility. diff --git a/tests/pipeline/model/training/augmentation/test_utils.py b/tests/pipeline/model/training/augmentation/test_utils.py index f792797..2f04d89 100644 --- a/tests/pipeline/model/training/augmentation/test_utils.py +++ b/tests/pipeline/model/training/augmentation/test_utils.py @@ -1,6 +1,7 @@ -from epochalyst.pipeline.model.training.augmentation import utils import torch +from epochalyst.pipeline.model.training.augmentation import utils + def set_torch_seed(seed: int = 42) -> None: """Set torch seed for reproducibility. diff --git a/tests/pipeline/model/training/models/test_timm.py b/tests/pipeline/model/training/models/test_timm.py index 43f7f00..8675d75 100644 --- a/tests/pipeline/model/training/models/test_timm.py +++ b/tests/pipeline/model/training/models/test_timm.py @@ -1,6 +1,7 @@ -from epochalyst.pipeline.model.training.models.timm import Timm import torch +from epochalyst.pipeline.model.training.models.timm import Timm + class TestTimm: diff --git a/tests/pipeline/model/training/test_pretrain_block.py b/tests/pipeline/model/training/test_pretrain_block.py index ecbe22d..9c4cdd4 100644 --- a/tests/pipeline/model/training/test_pretrain_block.py +++ b/tests/pipeline/model/training/test_pretrain_block.py @@ -1,4 +1,5 @@ import pytest + from epochalyst.pipeline.model.training.pretrain_block import PretrainBlock diff --git a/tests/pipeline/model/training/test_torch_trainer.py b/tests/pipeline/model/training/test_torch_trainer.py index a55b5a5..886caf3 100644 --- a/tests/pipeline/model/training/test_torch_trainer.py +++ b/tests/pipeline/model/training/test_torch_trainer.py @@ -1,19 +1,18 @@ import copy import functools -from dataclasses import dataclass import shutil +import time +from dataclasses import dataclass +from pathlib import Path from typing import Any from unittest.mock import patch -from pathlib import Path import numpy as np -import torch -from epochalyst.pipeline.model.training.torch_trainer import TorchTrainer import pytest +import torch -from tests.util import remove_cache_files - -TEMP_DIR = Path("tests/temp") +from epochalyst.pipeline.model.training.torch_trainer import TorchTrainer +from tests.constants import TEMP_DIR class TestTorchTrainer: @@ -57,22 +56,13 @@ def log_to_warning(self, message: str) -> None: pass @pytest.fixture(autouse=True) - def run_around_tests(self): - # Code that will run before each test - TEMP_DIR.mkdir(exist_ok=True) - - yield - - # Code that will run after each - if TEMP_DIR.exists(): - shutil.rmtree(TEMP_DIR) - + def run_always(self, setup_temp_dir): + pass def test_init_no_args(self): with pytest.raises(TypeError): TorchTrainer(n_folds=1) - def test_init_none_args(self): with pytest.raises(TypeError): TorchTrainer( @@ -83,7 +73,6 @@ def test_init_none_args(self): n_folds=1, ) - def test_init_proper_args(self): with pytest.raises(NotImplementedError): TorchTrainer( @@ -91,10 +80,9 @@ def test_init_proper_args(self): criterion=torch.nn.MSELoss(), optimizer=self.optimizer, n_folds=0, - model_name="Simple" + model_name="Simple", ) - def test_init_proper_args_with_implemented(self): tt = self.ImplementedTorchTrainer( model=self.simple_model, @@ -103,7 +91,6 @@ def test_init_proper_args_with_implemented(self): ) assert tt is not None - # Dataset concatenation def test__concat_datasets_in_order(self): tt = self.FullyImplementedTorchTrainer( @@ -130,7 +117,6 @@ def test__concat_datasets_in_order(self): assert (dataset[i][0] == x[i]).all() assert (dataset[i][1] == y[i]).all() - def test__concat_datasets_out_of_order(self): tt = self.FullyImplementedTorchTrainer( model=self.simple_model, @@ -158,7 +144,6 @@ def test__concat_datasets_out_of_order(self): assert (dataset[i][0] == x[i]).all() assert (dataset[i][1] == y[i]).all() - # Training def test_train_no_args(self): tt = self.ImplementedTorchTrainer( @@ -169,7 +154,6 @@ def test_train_no_args(self): with pytest.raises(TypeError): tt.train() - def test_train_no_train_indices(self): tt = self.ImplementedTorchTrainer( model=self.simple_model, @@ -179,7 +163,6 @@ def test_train_no_train_indices(self): with pytest.raises(ValueError): tt.train(torch.rand(10, 1), torch.rand(10)) - def test_train_no_test_indices(self): tt = self.ImplementedTorchTrainer( model=self.simple_model, @@ -193,32 +176,29 @@ def test_train_no_test_indices(self): train_indices=[0, 1, 2, 3, 4, 5, 6, 7], ) - def test_train(self): tt = self.FullyImplementedTorchTrainer( model=self.simple_model, criterion=torch.nn.MSELoss(), optimizer=self.optimizer, ) - + x = torch.rand(10, 1) y = torch.rand(10) tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[8, 9]) - def test_train_trained(self): tt = self.FullyImplementedTorchTrainer( model=self.simple_model, criterion=torch.nn.MSELoss(), optimizer=self.optimizer, ) - + x = torch.rand(10, 1) y = torch.rand(10) tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[8, 9]) tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[8, 9]) - def test_train_full(self): tt = self.FullyImplementedTorchTrainer( model=self.simple_model, @@ -226,14 +206,13 @@ def test_train_full(self): optimizer=self.optimizer, ) tt.n_folds = 0 - + x = torch.rand(10, 1) y = torch.rand(10) tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], test_indices=[]) tt.predict(x) - def test_early_stopping(self): tt = self.FullyImplementedTorchTrainer( model=self.simple_model, @@ -241,12 +220,11 @@ def test_early_stopping(self): optimizer=self.optimizer, patience=-1, ) - + x = torch.rand(10, 1) y = torch.rand(10) tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], test_indices=[]) - # Test predict def test_predict_no_args(self): tt = self.ImplementedTorchTrainer( @@ -263,7 +241,7 @@ def test_predict(self): criterion=torch.nn.MSELoss(), optimizer=self.optimizer, ) - + x = torch.rand(10, 1) y = torch.rand(10) tt.train( @@ -271,16 +249,15 @@ def test_predict(self): ) tt.predict(x) - def test_predict_3fold(self): tt = self.FullyImplementedTorchTrainer( model=self.simple_model, criterion=torch.nn.MSELoss(), optimizer=self.optimizer, ) - + tt.n_folds = 3 - + x = torch.rand(10, 1) y = torch.rand(10) tt.train( @@ -294,40 +271,38 @@ def test_predict_3fold(self): ) tt.predict(x) - def test_predict_train_full(self): tt = self.FullyImplementedTorchTrainer( model=self.simple_model, criterion=torch.nn.MSELoss(), optimizer=self.optimizer, ) - + tt.n_folds = 0 - + x = torch.rand(10, 1) y = torch.rand(10) tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[]) tt.predict(x) - def test_predict2(self): tt = self.FullyImplementedTorchTrainer( model=self.simple_model, criterion=torch.nn.MSELoss(), optimizer=self.optimizer, ) - + x = torch.rand(10, 1) y = torch.rand(10) tt.train( x, y, train_indices=np.array([0, 1, 2, 3, 4, 5, 6, 7]), - test_indices=np.array([8, 9]), fold=0 + test_indices=np.array([8, 9]), + fold=0, ) tt.predict(x) - def test_predict_all(self): tt = self.FullyImplementedTorchTrainer( model=self.simple_model, @@ -335,19 +310,19 @@ def test_predict_all(self): optimizer=self.optimizer, to_predict="all", ) - + x = torch.rand(10, 1) y = torch.rand(10) train_preds = tt.train( x, y, train_indices=np.array([0, 1, 2, 3, 4, 5, 6, 7]), - test_indices=np.array([8, 9]), fold=0 + test_indices=np.array([8, 9]), + fold=0, ) preds = tt.predict(x) assert len(train_preds[0]) == 10 assert len(preds) == 10 - def test_predict_2d(self): tt = self.FullyImplementedTorchTrainer( @@ -356,20 +331,20 @@ def test_predict_2d(self): optimizer=self.optimizer, to_predict="all", ) - + x = torch.rand(10, 2) y = torch.rand(10, 2) train_preds = tt.train( x, y, train_indices=np.array([0, 1, 2, 3, 4, 5, 6, 7]), - test_indices=np.array([8, 9]), fold=0 + test_indices=np.array([8, 9]), + fold=0, ) preds = tt.predict(x) assert len(train_preds[0]) == 10 assert preds.shape == (10, 2) assert len(preds) == 10 - def test_predict_partial(self): tt = self.FullyImplementedTorchTrainer( @@ -378,20 +353,20 @@ def test_predict_partial(self): optimizer=self.optimizer, to_predict="test", ) - + x = torch.rand(10, 1) y = torch.rand(10) train_preds = tt.train( x, y, train_indices=np.array([0, 1, 2, 3, 4, 5, 6, 7]), - test_indices=np.array([8, 9]), fold=0 + test_indices=np.array([8, 9]), + fold=0, ) preds = tt.predict(x) assert len(train_preds[0]) == 2 assert len(preds) == 10 - def test_predict_none(self): tt = self.FullyImplementedTorchTrainer( model=self.simple_model, @@ -399,20 +374,20 @@ def test_predict_none(self): optimizer=self.optimizer, to_predict="none", ) - + x = torch.rand(10, 1) y = torch.rand(10) train_preds = tt.train( x, y, train_indices=np.array([0, 1, 2, 3, 4, 5, 6, 7]), - test_indices=np.array([8, 9]), fold=0 + test_indices=np.array([8, 9]), + fold=0, ) preds = tt.predict(x) assert len(train_preds[0]) == 10 assert len(preds) == 10 - def test_predict_no_model_trained(self): tt = self.FullyImplementedTorchTrainer( model=self.simple_model, @@ -422,7 +397,6 @@ def test_predict_no_model_trained(self): with pytest.raises(FileNotFoundError): tt.predict(torch.rand(10, 1)) - # Test with scheduler def test_train_with_scheduler(self): tt = self.FullyImplementedTorchTrainer( @@ -431,12 +405,11 @@ def test_train_with_scheduler(self): optimizer=self.optimizer, scheduler=self.scheduler, ) - + x = torch.rand(10, 1) y = torch.rand(10) tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[8, 9]) - # Test 1 gpu training def test_train_one_gpu(self): with patch("torch.cuda.device_count", return_value=1): @@ -449,7 +422,6 @@ def test_train_one_gpu(self): x = torch.rand(10, 1) y = torch.rand(10) tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[8, 9]) - def test_train_one_gpu_saved(self): with patch("torch.cuda.device_count", return_value=1): @@ -464,7 +436,6 @@ def test_train_one_gpu_saved(self): tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[8, 9]) tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[8, 9]) - def test_train_two_gpu_saved(self): # If test is run on a machine with 2 or more GPUs, this test will run else it will be skipped if torch.cuda.device_count() < 2: @@ -482,7 +453,6 @@ def test_train_two_gpu_saved(self): tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[8, 9]) tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[8, 9]) - def test_early_stopping_no_patience(self): tt = self.FullyImplementedTorchTrainer( model=self.simple_model, @@ -511,7 +481,6 @@ def test_early_stopping_no_patience(self): # Assert model chnages after training assert tt.model.state_dict != orig_state_dict - def test_checkpointing(self): tt = self.FullyImplementedTorchTrainer( model=self.simple_model, @@ -520,12 +489,40 @@ def test_checkpointing(self): patience=-1, checkpointing_enabled=True, checkpointing_keep_every=1, + checkpointing_resume_if_exists=True, ) x = torch.rand(10, 1) y = torch.rand(10) + + # Train once + time_temp = time.time() + tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[8, 9]) + spent_time_first_run = time.time() - time_temp + + # Check if checkpoints exist + saved_checkpoints = list(tt.trained_models_directory.glob(f"*_checkpoint_*.pt")) + assert len(saved_checkpoints) == tt.epochs + + # Remove model and all but the 2nd to last checkpoint + start_epoch = tt.epochs - 2 + epochs = [ + int(checkpoint.stem.split("_")[-1]) for checkpoint in saved_checkpoints + ] + checkpoint_to_keep = saved_checkpoints[epochs.index(start_epoch)] + print(checkpoint_to_keep) + for file in tt.trained_models_directory.glob("*.pt"): + if file != checkpoint_to_keep: + file.unlink() + + # Train again + time_temp = time.time() tt.train(x, y, train_indices=[0, 1, 2, 3, 4, 5, 6, 7], test_indices=[8, 9]) + spent_time_second_run = time.time() - time_temp # Check if checkpoints exist saved_checkpoints = list(tt.trained_models_directory.glob(f"*_checkpoint_*.pt")) - assert len(saved_checkpoints) == tt.epochs \ No newline at end of file + assert len(saved_checkpoints) == tt.epochs - start_epoch + + # Check if training time was signficicantly less the second time + assert spent_time_second_run < (spent_time_first_run/2) diff --git a/tests/pipeline/model/training/test_training.py b/tests/pipeline/model/training/test_training.py index 85840ad..3891b2a 100644 --- a/tests/pipeline/model/training/test_training.py +++ b/tests/pipeline/model/training/test_training.py @@ -1,11 +1,15 @@ -from epochalyst.pipeline.model.training.training import TrainingPipeline -from agogos.training import Trainer +from pathlib import Path + import numpy as np +import pytest +from agogos.training import Trainer + +from epochalyst.pipeline.model.training.training import TrainingPipeline from epochalyst.pipeline.model.training.training_block import TrainingBlock -from tests.util import remove_cache_files +from tests.constants import TEMP_DIR -class TestTrainingBlock(TrainingBlock): +class ExampleTrainingBlock(TrainingBlock): def custom_train(self, x, y, **train_args): if x is None: return None, y @@ -23,7 +27,7 @@ def log_to_terminal(self, message): pass -class CustomTrainingPipeline(TrainingPipeline): +class ExampleTrainingPipeline(TrainingPipeline): def log_to_debug(self, message: str) -> None: return None @@ -32,6 +36,12 @@ def log_to_terminal(self, message: str) -> None: class TestTrainingPipeline: + cache_path = TEMP_DIR + + @pytest.fixture(autouse=True) + def run_always(self, setup_temp_dir): + pass + def test_training_pipeline_init(self): tp = TrainingPipeline() assert tp is not None @@ -47,22 +57,22 @@ def test_training_pipeline_predict(self): assert tp.predict(None) is None def test_training_pipeline_with_steps(self): - t1 = TestTrainingBlock() - t2 = TestTrainingBlock() - tp = CustomTrainingPipeline(steps=[t1, t2]) + t1 = ExampleTrainingBlock() + t2 = ExampleTrainingBlock() + tp = ExampleTrainingPipeline(steps=[t1, t2]) assert tp.train(None, None) == (None, None) def test_training_pipeline_with_cache(self): - t1 = TestTrainingBlock() - t2 = TestTrainingBlock() + t1 = ExampleTrainingBlock() + t2 = ExampleTrainingBlock() - tp = CustomTrainingPipeline(steps=[t1, t2]) + tp = ExampleTrainingPipeline(steps=[t1, t2]) cache_args = { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } x, y = tp.train(1, 1, cache_args=cache_args) @@ -73,21 +83,20 @@ def test_training_pipeline_with_cache(self): pred = tp.predict(1, cache_args=cache_args) new_pred = tp.predict(1, cache_args=cache_args) assert pred == new_pred - remove_cache_files() def test_training_pipeline_with_halfway_cache(self): - t1 = TestTrainingBlock() - t2 = TestTrainingBlock() + t1 = ExampleTrainingBlock() + t2 = ExampleTrainingBlock() - tp1 = CustomTrainingPipeline(steps=[t1]) - tp2 = CustomTrainingPipeline(steps=[t1, t2]) + tp1 = ExampleTrainingPipeline(steps=[t1]) + tp2 = ExampleTrainingPipeline(steps=[t1, t2]) training_args = { - "TestTrainingBlock": { + "ExampleTrainingBlock": { "cache_args": { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } } } @@ -103,21 +112,20 @@ def test_training_pipeline_with_halfway_cache(self): assert tp1.predict(np.array([1]), **training_args) == np.array([2]) assert tp2.predict(np.array([3]), **training_args) == np.array([4]) - remove_cache_files() def test_training_pipeline_with_halfway_cache_no_step_cache_args(self): - t1 = TestTrainingBlock() - t2 = TestTrainingBlock() + t1 = ExampleTrainingBlock() + t2 = ExampleTrainingBlock() - tp1 = CustomTrainingPipeline(steps=[t1]) - tp2 = CustomTrainingPipeline(steps=[t1, t2]) + tp1 = ExampleTrainingPipeline(steps=[t1]) + tp2 = ExampleTrainingPipeline(steps=[t1, t2]) training_args = { - "TestTrainingBlock": { + "ExampleTrainingBlock": { "wrong_args": { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } } } @@ -133,7 +141,6 @@ def test_training_pipeline_with_halfway_cache_no_step_cache_args(self): assert tp1.predict(np.array([1]), **training_args) == np.array([2]) assert tp2.predict(np.array([3]), **training_args) != np.array([4]) - remove_cache_files() def test_training_pipeline_with_halfway_cache_not_instance_cacher(self): class ImplementedTrainer(Trainer): @@ -144,22 +151,22 @@ def predict(self, x, **pred_args): return x * 2 t1 = ImplementedTrainer() - t2 = TestTrainingBlock() + t2 = ExampleTrainingBlock() - tp1 = CustomTrainingPipeline(steps=[t1]) - tp2 = CustomTrainingPipeline(steps=[t1, t2]) + tp1 = ExampleTrainingPipeline(steps=[t1]) + tp2 = ExampleTrainingPipeline(steps=[t1, t2]) training_args = { - "TestTrainingBlock": { + "ExampleTrainingBlock": { "cache_args": { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } } } - assert tp1.train(np.array([1]), np.array([1]), **training_args) == ( + assert tp1.train(np.array([1]), np.array([1])) == ( np.array([2]), np.array([1]), ) @@ -170,4 +177,3 @@ def predict(self, x, **pred_args): assert tp1.predict(np.array([1]), **training_args) == np.array([2]) assert tp2.predict(np.array([2]), **training_args) == np.array([8]) - remove_cache_files() diff --git a/tests/pipeline/model/training/test_training_block.py b/tests/pipeline/model/training/test_training_block.py index 4d42be2..15d7a35 100644 --- a/tests/pipeline/model/training/test_training_block.py +++ b/tests/pipeline/model/training/test_training_block.py @@ -1,9 +1,16 @@ +import shutil +from pathlib import Path + import pytest + from epochalyst.pipeline.model.training.training_block import TrainingBlock -from tests.util import remove_cache_files + +TEMP_DIR = Path("tests/temp") class TestTrainingBlock: + cache_path = TEMP_DIR + def test_training_block_init(self): tb = TrainingBlock() assert tb is not None @@ -30,9 +37,7 @@ def custom_predict(self, x: int) -> int: assert tb.train(1, 1) == (1, 1) assert tb.predict(1) == 1 - remove_cache_files() - - def test_training_block_caching(self): + def test_training_block_caching(self, setup_temp_dir): class TestTrainingBlockImpl(TrainingBlock): def custom_train(self, x: int, y: int) -> tuple[int, int]: return x, y @@ -50,7 +55,7 @@ def log_to_terminal(self, message: str) -> None: cache_args = { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } x, y = tb.train(1, 1, cache_args=cache_args) @@ -61,5 +66,3 @@ def log_to_terminal(self, message: str) -> None: pred = tb.predict(1, cache_args=cache_args) new_pred = tb.predict(1, cache_args=cache_args) assert pred == new_pred - - remove_cache_files() diff --git a/tests/pipeline/model/transformation/test_transformation.py b/tests/pipeline/model/transformation/test_transformation.py index 196101d..e9444f8 100644 --- a/tests/pipeline/model/transformation/test_transformation.py +++ b/tests/pipeline/model/transformation/test_transformation.py @@ -1,15 +1,18 @@ -from epochalyst.pipeline.model.transformation.transformation import ( - TransformationPipeline, -) -from epochalyst.pipeline.model.transformation.transformation_block import ( - TransformationBlock, -) -from agogos.transforming import Transformer +import shutil +from pathlib import Path + import numpy as np -from tests.util import remove_cache_files +import pytest +from agogos.transforming import Transformer + +from epochalyst.pipeline.model.transformation.transformation import \ + TransformationPipeline +from epochalyst.pipeline.model.transformation.transformation_block import \ + TransformationBlock +from tests.constants import TEMP_DIR -class TestTransformationBlock(TransformationBlock): +class ExampleTransformationBlock(TransformationBlock): def log_to_debug(self, message): pass @@ -31,6 +34,12 @@ def log_to_terminal(self, message: str) -> None: class TestTransformationPipeline: + cache_path = TEMP_DIR + + @pytest.fixture(autouse=True) + def run_always(self, setup_temp_dir): + pass + def test_transformation_pipeline_init(self): tp = CustomTransformationPipeline() assert tp.steps is not None @@ -43,91 +52,87 @@ def test_transformation_pipeline_transform(self): assert tp.transform(x, transform_args={"a": 1, "b": 2}) == x def test_transformation_pipeline_with_steps(self): - t1 = TestTransformationBlock() - t2 = TestTransformationBlock() + t1 = ExampleTransformationBlock() + t2 = ExampleTransformationBlock() tp = CustomTransformationPipeline(steps=[t1, t2]) assert tp.transform(None) is None def test_transformation_pipeline_with_cache(self): - t1 = TestTransformationBlock() - t2 = TestTransformationBlock() + t1 = ExampleTransformationBlock() + t2 = ExampleTransformationBlock() tp = CustomTransformationPipeline(steps=[t1, t2]) cache_args = { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } assert tp.transform(np.array([1]), cache_args=cache_args) == np.array([4]) assert tp.transform(np.array([1]), cache_args=cache_args) == np.array([4]) - remove_cache_files() def test_transformation_pipeline_with_halfway_cache(self): - t1 = TestTransformationBlock() - t2 = TestTransformationBlock() + t1 = ExampleTransformationBlock() + t2 = ExampleTransformationBlock() tp1 = CustomTransformationPipeline(steps=[t1, t2]) tp2 = CustomTransformationPipeline(steps=[t1, t2]) transform_args = { - "TestTransformationBlock": { + "ExampleTransformationBlock": { "cache_args": { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } } } assert tp1.transform(np.array([1]), **transform_args) == np.array([4]) assert tp2.transform(np.array([1]), **transform_args) == np.array([4]) - remove_cache_files() def test_transformation_pipeline_with_halfway_cache_no_step_cache_args(self): - t1 = TestTransformationBlock() - t2 = TestTransformationBlock() + t1 = ExampleTransformationBlock() + t2 = ExampleTransformationBlock() tp1 = CustomTransformationPipeline(steps=[t1, t2]) tp2 = CustomTransformationPipeline(steps=[t1, t2]) transform_args = { - "TestTransformationBlock": { + "ExampleTransformationBlock": { "wrong_args": { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } } } assert tp1.transform(np.array([1]), **transform_args) == np.array([4]) assert tp2.transform(np.array([1]), **transform_args) == np.array([4]) - remove_cache_files() def test_transformation_pipeline_with_halfway_cache_not_instance_cacher(self): class ImplementedTransformer(Transformer): def transform(self, x, **transform_args): return x * 2 - t1 = TestTransformationBlock() + t1 = ExampleTransformationBlock() t2 = ImplementedTransformer() tp1 = CustomTransformationPipeline(steps=[t1, t2]) tp2 = CustomTransformationPipeline(steps=[t1, t2]) transform_args = { - "TestTransformationBlock": { + "ExampleTransformationBlock": { "cache_args": { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } } } assert tp1.transform(np.array([1]), **transform_args) == np.array([4]) assert tp2.transform(np.array([2]), **transform_args) == np.array([4]) - remove_cache_files() diff --git a/tests/pipeline/model/transformation/test_transformation_block.py b/tests/pipeline/model/transformation/test_transformation_block.py index 7a3cb2e..0c04245 100644 --- a/tests/pipeline/model/transformation/test_transformation_block.py +++ b/tests/pipeline/model/transformation/test_transformation_block.py @@ -1,12 +1,18 @@ +import shutil +from pathlib import Path + import numpy as np -from tests.util import remove_cache_files import pytest -from epochalyst.pipeline.model.transformation.transformation_block import ( - TransformationBlock, -) + +from epochalyst.pipeline.model.transformation.transformation_block import \ + TransformationBlock + +TEMP_DIR = Path("tests/temp") class TestTransformationBlock: + cache_path = TEMP_DIR + def test_transformation_block_init(self): tb = TransformationBlock() assert tb is not None @@ -55,7 +61,7 @@ def custom_transform(self, data: int, **transform_args) -> int: assert tb.transform(1) == 2 assert tb.transform(1) == 2 - def test_tb_custom_transform_implementation_with_cache(self): + def test_tb_custom_transform_implementation_with_cache(self, setup_temp_dir): class TestTransformationBlockImpl(TransformationBlock): def custom_transform(self, data: np.ndarray[int], **transform_args) -> int: return data * 2 @@ -70,10 +76,8 @@ def log_to_terminal(self, message: str) -> None: cache_args = { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } assert tb.transform(np.array([1]), cache_args=cache_args) == np.array([2]) assert tb.transform(np.array([1]), cache_args=cache_args) == np.array([2]) - - remove_cache_files() diff --git a/tests/pipeline/test_ensemble.py b/tests/pipeline/test_ensemble.py index 03f2511..42c7b55 100644 --- a/tests/pipeline/test_ensemble.py +++ b/tests/pipeline/test_ensemble.py @@ -1,17 +1,20 @@ -from epochalyst.pipeline.ensemble import EnsemblePipeline -from tests.util import remove_cache_files +import shutil +from pathlib import Path +from typing import Any + import numpy as np +import pytest + +from epochalyst.pipeline.ensemble import EnsemblePipeline from epochalyst.pipeline.model.model import ModelPipeline -from epochalyst.pipeline.model.transformation.transformation import ( - TransformationPipeline, -) -from epochalyst.pipeline.model.transformation.transformation_block import ( - TransformationBlock, -) -from typing import Any +from epochalyst.pipeline.model.transformation.transformation import \ + TransformationPipeline +from epochalyst.pipeline.model.transformation.transformation_block import \ + TransformationBlock +from tests.constants import TEMP_DIR -class TestTransformation(TransformationBlock): +class ExampleTransformation(TransformationBlock): def log_to_debug(self, message: str) -> None: pass @@ -19,7 +22,7 @@ def custom_transform(self, data: Any, **transform_args: Any) -> Any: return data * 2 -class TestTransformationPipeline(TransformationPipeline): +class ExampleTransformationPipeline(TransformationPipeline): def log_to_debug(self, message: str) -> None: pass @@ -28,6 +31,12 @@ def log_to_terminal(self, message: str) -> None: class TestEnsemble: + cache_path = TEMP_DIR + + @pytest.fixture(autouse=True) + def run_always(self, setup_temp_dir): + pass + def test_ensemble_pipeline_init(self): ensemble = EnsemblePipeline() assert ensemble is not None @@ -72,7 +81,7 @@ def test_ensemble_get_cache_exists_false(self): cache_args = { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } assert ensemble.get_x_cache_exists(cache_args) is False @@ -87,16 +96,16 @@ def test_ensemble_get_cache_exists_models_false(self): cache_args = { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } assert ensemble.get_x_cache_exists(cache_args) is False assert ensemble.get_y_cache_exists(cache_args) is False def test_ensemble_get_cache_exists_true(self): - test_transformation = TestTransformation() - x_sys = TestTransformationPipeline([test_transformation]) - y_sys = TestTransformationPipeline([test_transformation]) + test_transformation = ExampleTransformation() + x_sys = ExampleTransformationPipeline([test_transformation]) + y_sys = ExampleTransformationPipeline([test_transformation]) model1 = ModelPipeline(x_sys=x_sys, y_sys=y_sys) model2 = ModelPipeline(x_sys=x_sys, y_sys=y_sys) @@ -106,7 +115,7 @@ def test_ensemble_get_cache_exists_true(self): cache_args = { "output_data_type": "numpy_array", "storage_type": ".npy", - "storage_path": "tests/cache", + "storage_path": f"{self.cache_path}", } train_args = { @@ -126,5 +135,3 @@ def test_ensemble_get_cache_exists_true(self): assert ensemble.get_x_cache_exists(cache_args) is True assert ensemble.get_y_cache_exists(cache_args) is True - - remove_cache_files() diff --git a/tests/util.py b/tests/util.py deleted file mode 100644 index 93eb658..0000000 --- a/tests/util.py +++ /dev/null @@ -1,18 +0,0 @@ -import os -import shutil - - -def remove_cache_files() -> None: - """Remove all files and directories in the given path.""" - path = "tests/cache" - for filename in os.listdir(path): - if filename in ["README", "README.md"]: - continue - file_path = os.path.join(path, filename) - try: - if os.path.isfile(file_path) or os.path.islink(file_path): - os.unlink(file_path) - elif os.path.isdir(file_path): - shutil.rmtree(file_path) - except Exception as e: - print(f"Failed to delete {file_path}. Reason: {e}")