Coverage for tests/test_model_tuning.py: 94%

1import unittest 

2import pandas as pd 

3from addmo.s3_model_tuning.models.model_factory import ModelFactory 

4from addmo.s3_model_tuning.config.model_tuning_config import ModelTuningExperimentConfig 

5from addmo.s3_model_tuning.config.model_tuning_config import ModelTunerConfig 

6from addmo.s3_model_tuning.model_tuner import ModelTuner 

7from addmo.util.load_save import load_data 

8from addmo.util.load_save import load_config_from_json 

9from addmo.util.data_handling import split_target_features 

10 

11 

12class TestModelTuner(unittest.TestCase): 

13 """ 

14 Unit tests for ModelTuner using pre-defined config. 

15 """ 

16 def setUp(self): 

17 self.config = ModelTuningExperimentConfig() 

18 

19 def test_model_tuning(self): 

20 model_tuner = ModelTuner(config=self.config._config_model_tuner) 

21 

22 # Load the system_data 

23 xy_tuned = load_data(self.config.abs_path_to_data) 

24 

25 # Select training and validation period 

26 if hasattr(self.config, "start_train_val") and hasattr(self.config, "stop_train_val"): 

27 xy_tuned_train_val = xy_tuned.loc[self.config.start_train_val: self.config.stop_train_val] 

28 else: 

29 xy_tuned_train_val = xy_tuned 

30 

31 x_train_val, y_train_val = split_target_features(self.config.name_of_target, xy_tuned_train_val) 

32 

33 # Tune the models 

34 model_dict = model_tuner.tune_all_models(x_train_val, y_train_val) 

35 

36 # Get the best model 

37 best_model_name = model_tuner.get_best_model_name(model_dict) 

38 best_model = model_tuner.get_model(model_dict, best_model_name) 

39 

40 # Validation 

41 self.assertFalse(xy_tuned.empty, "loaded data is empty") 

42 self.assertIsInstance(xy_tuned, pd.DataFrame, "load_data() not returning a pd.DataFrame") 

43 self.assertIsInstance(xy_tuned_train_val, pd.DataFrame,"train_val data is no longer a pd.DataFrame") 

44 self.assertFalse(xy_tuned_train_val.empty, "train_val data is empty") 

45 self.assertIsInstance(model_dict, dict, "tune_all_models() should return a dict") 

46 self.assertSetEqual(set(model_dict.keys()), set(model_tuner.config.models), "Returned models do not match expected models") 

47 

48 for model_name, model in model_dict.items(): 

49 self.assertIsNotNone(model, f"Model {model_name} cannot be None") 

50 expected_model_class = ModelFactory.model_factory(model_name).__class__ 

51 self.assertIsInstance(model, expected_model_class, f"Model should be {expected_model_class} but got {type(model)}") 

52 

53 self.assertIsNotNone(best_model, "best_model is None") 

54 

55if __name__ == "__main__": 

56 unittest.main()