SIENTIAPDE-1430: Add unit tests for _configure_datetime_index (handling various datetime column formats), _init_data_preprocessor (with removed intervals), and _extract_model_equation (supporting polynomial features) in TrainingRepository.

tests/utils/repository/test_training_repository.py

@@ -791,3 +791,228 @@ class TestTrain:
         assert hasattr(result, 'regr')
         assert hasattr(result, 'scaler_dict')
         assert result.params == sample_params
+
+
+# ============================================================================
+# Tests for _configure_datetime_index
+# ============================================================================
+
+
+class TestConfigureDatetimeIndex:
+    """Tests for _configure_datetime_index method."""
+
+    @pytest.fixture
+    def training_repo(self, mock_logger):
+        """Create a TrainingRepository instance."""
+        return TrainingRepository(logger=mock_logger)
+
+    def test_configure_datetime_index_already_datetime(self, training_repo):
+        """Test _configure_datetime_index when index is already DatetimeIndex."""
+        data = pd.DataFrame(
+            {'var1': [1, 2, 3], 'var2': [4, 5, 6]},
+            index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']),
+        )
+
+        result = training_repo._configure_datetime_index(data)
+
+        assert isinstance(result.index, pd.DatetimeIndex)
+        assert len(result) == 3
+
+    def test_configure_datetime_index_with_timestamp_column(self, training_repo):
+        """Test _configure_datetime_index with 'timestamp' column."""
+        data = pd.DataFrame(
+            {
+                'timestamp': ['2023-01-01', '2023-01-02', '2023-01-03'],
+                'var1': [1, 2, 3],
+                'var2': [4, 5, 6],
+            }
+        )
+
+        result = training_repo._configure_datetime_index(data)
+
+        assert isinstance(result.index, pd.DatetimeIndex)
+        assert 'timestamp' not in result.columns
+
+    def test_configure_datetime_index_with_date_column(self, training_repo):
+        """Test _configure_datetime_index with 'date' column."""
+        data = pd.DataFrame(
+            {
+                'date': ['2023-01-01', '2023-01-02', '2023-01-03'],
+                'var1': [1, 2, 3],
+            }
+        )
+
+        result = training_repo._configure_datetime_index(data)
+
+        assert isinstance(result.index, pd.DatetimeIndex)
+        assert 'date' not in result.columns
+
+    def test_configure_datetime_index_with_datetime_column(self, training_repo):
+        """Test _configure_datetime_index with 'datetime' column."""
+        data = pd.DataFrame(
+            {
+                'datetime': ['2023-01-01', '2023-01-02', '2023-01-03'],
+                'var1': [1, 2, 3],
+            }
+        )
+
+        result = training_repo._configure_datetime_index(data)
+
+        assert isinstance(result.index, pd.DatetimeIndex)
+        assert 'datetime' not in result.columns
+
+    def test_configure_datetime_index_first_column_datetime(self, training_repo):
+        """Test _configure_datetime_index when first column looks like datetime."""
+        data = pd.DataFrame(
+            {
+                'my_date': ['2023-01-01', '2023-01-02', '2023-01-03'],
+                'var1': [1, 2, 3],
+            }
+        )
+
+        result = training_repo._configure_datetime_index(data)
+
+        assert isinstance(result.index, pd.DatetimeIndex)
+        assert 'my_date' not in result.columns
+
+    @pytest.mark.filterwarnings('ignore::UserWarning')
+    def test_configure_datetime_index_no_timestamp_column(self, training_repo):
+        """Test _configure_datetime_index when no timestamp column found."""
+        data = pd.DataFrame(
+            {
+                'var1': ['text_a', 'text_b', 'text_c'],
+                'var2': ['text_d', 'text_e', 'text_f'],
+            }
+        )
+
+        result = training_repo._configure_datetime_index(data)
+
+        # Should return original data unchanged (no valid datetime columns)
+        assert 'var1' in result.columns
+        assert 'var2' in result.columns
+
+    @pytest.mark.filterwarnings('ignore::UserWarning')
+    def test_configure_datetime_index_invalid_timestamp_column(self, training_repo):
+        """Test _configure_datetime_index with invalid timestamp values."""
+        data = pd.DataFrame(
+            {
+                'timestamp': ['not_a_date', 'also_not', 'nope'],
+                'var1': [1, 2, 3],
+            }
+        )
+
+        result = training_repo._configure_datetime_index(data)
+
+        # Should skip invalid column and try first column
+        assert 'var1' in result.columns
+
+    @pytest.mark.filterwarnings('ignore::UserWarning')
+    def test_configure_datetime_index_invalid_first_column(self, training_repo):
+        """Test _configure_datetime_index when first column is not datetime."""
+        data = pd.DataFrame(
+            {
+                'var1': ['a', 'b', 'c'],
+                'var2': [1, 2, 3],
+            }
+        )
+
+        result = training_repo._configure_datetime_index(data)
+
+        # Should return original data unchanged
+        assert 'var1' in result.columns
+        assert 'var2' in result.columns
+
+    def test_configure_datetime_index_first_column_all_nan(self, training_repo):
+        """Test _configure_datetime_index when first column has all NaN values."""
+        data = pd.DataFrame(
+            {
+                'first_col': [np.nan, np.nan, np.nan],
+                'var1': [1, 2, 3],
+            }
+        )
+
+        result = training_repo._configure_datetime_index(data)
+
+        # Should return original data unchanged (first column has no valid values)
+        assert 'first_col' in result.columns
+        assert 'var1' in result.columns
+
+
+# ============================================================================
+# Tests for _init_data_preprocessor with removed_intervals
+# ============================================================================
+
+
+class TestInitDataPreprocessorWithRemovedIntervals:
+    """Tests for _init_data_preprocessor with removed_intervals."""
+
+    @pytest.fixture
+    def training_repo(self, mock_logger):
+        """Create a TrainingRepository instance."""
+        return TrainingRepository(logger=mock_logger)
+
+    def test_init_data_preprocessor_with_removed_intervals(self, training_repo, sample_params):
+        """Test _init_data_preprocessor with removed_intervals."""
+        sample_params.removed_intervals = [
+            ['2023-01-01', '2023-01-02'],
+            ['2023-02-01', '2023-02-02'],
+        ]
+
+        preprocessor = training_repo._init_data_preprocessor(sample_params)
+
+        assert preprocessor is not None
+        assert preprocessor.removed_intervals is not None
+        assert len(preprocessor.removed_intervals) == 2
+
+    def test_init_data_preprocessor_with_tuple_intervals(self, training_repo, sample_params):
+        """Test _init_data_preprocessor with tuple intervals."""
+        sample_params.removed_intervals = [('2023-01-01', '2023-01-02')]
+
+        preprocessor = training_repo._init_data_preprocessor(sample_params)
+
+        assert preprocessor is not None
+        assert preprocessor.removed_intervals is not None
+
+
+# ============================================================================
+# Tests for _extract_model_equation with polynomial features
+# ============================================================================
+
+
+class TestExtractModelEquationPolynomial:
+    """Tests for _extract_model_equation with polynomial features."""
+
+    @pytest.fixture
+    def training_repo(self, mock_logger):
+        """Create a TrainingRepository instance."""
+        return TrainingRepository(logger=mock_logger)
+
+    def test_extract_model_equation_polynomial(self, training_repo, sample_params):
+        """Test _extract_model_equation with polynomial features."""
+        # Create a mock regressor with polynomial features
+        mock_regr = MagicMock()
+        mock_regr.regr.coef_ = np.array([0.5, 0.3, 0.2])
+        mock_regr.regr.intercept_ = 1.0
+        mock_regr.poly_feature_names = ['var1', 'var2', 'var1^2']
+
+        sample_params.degree = 2
+
+        result = training_repo._extract_model_equation(mock_regr, sample_params)
+
+        assert 'equation_string' in result
+        assert 'latex_equation' in result
+        assert 'var1' in result['equation_string']
+
+    def test_extract_model_equation_linear(self, training_repo, sample_params):
+        """Test _extract_model_equation with linear features."""
+        mock_regr = MagicMock()
+        mock_regr.regr.coef_ = np.array([0.5, 0.3])
+        mock_regr.regr.intercept_ = 1.0
+        mock_regr.poly_feature_names = None
+
+        sample_params.degree = 1
+
+        result = training_repo._extract_model_equation(mock_regr, sample_params)
+
+        assert 'equation_string' in result
+        assert 'latex_equation' in result