Modify and complete test_forecast_ar.py and test_clean_apple.py so that all functions are tested and number of assert statements are reduced.

Author: Lenr4

src/lennart_epp/data_management/clean_apple.py

@@ -1,7 +1,7 @@
 import pandas as pd
 
 
-def _select_and_rename_column(df: pd.DataFrame) -> pd.DataFrame:
+def _rename_column(df: pd.DataFrame) -> pd.DataFrame:
     """Select the 'Close' column and rename it to 'close_price'.
 
     Args:
@@ -92,7 +92,7 @@ def clean_apple_data(df: pd.DataFrame) -> pd.DataFrame:
     _validate_dataframe(df)
 
     df = _convert_to_datetime(df)
-    df = _select_and_rename_column(df)
+    df = _rename_column(df)
     df = _handle_missing_values(df)
     df = _remove_duplicates(df)
     df = _convert_to_numeric(df)

tests/analysis/test_forecast_ar.py

@@ -19,16 +19,45 @@ def test_integrated_coefficients():
     return pd.DataFrame({"coefficient": coeffs})
 
 
-def test_forecast_ar_multi_step(test_dataframe, test_integrated_coefficients):
+def test_forecast_ar_multi_step_output_type(
+    test_dataframe, test_integrated_coefficients
+):
+    """Test that forecast_ar_multi_step returns a pandas Series."""
     forecast_steps = 10
     forecasts = forecast_ar_multi_step(
         test_dataframe, test_integrated_coefficients, forecast_steps
     )
 
     assert isinstance(forecasts, pd.Series)
+
+
+def test_forecast_ar_multi_step_length(test_dataframe, test_integrated_coefficients):
+    """Test that the forecast length matches the requested number of steps."""
+    forecast_steps = 10
+    forecasts = forecast_ar_multi_step(
+        test_dataframe, test_integrated_coefficients, forecast_steps
+    )
+
     assert len(forecasts) == forecast_steps
+
+
+def test_forecast_ar_multi_step_no_nans(test_dataframe, test_integrated_coefficients):
+    """Test that the forecast does not contain NaN values."""
+    forecast_steps = 10
+    forecasts = forecast_ar_multi_step(
+        test_dataframe, test_integrated_coefficients, forecast_steps
+    )
+
     assert not forecasts.isna().any()
 
+
+def test_forecast_ar_multi_step_index(test_dataframe, test_integrated_coefficients):
+    """Test that the forecast index matches the expected date range."""
+    forecast_steps = 10
+    forecasts = forecast_ar_multi_step(
+        test_dataframe, test_integrated_coefficients, forecast_steps
+    )
+
     expected_index_start = test_dataframe.index[-344]
     expected_index = pd.date_range(
         start=expected_index_start, periods=forecast_steps, freq="D"

tests/data_management/test_clean_apple.py

@@ -7,8 +7,9 @@
     _convert_to_numeric,
     _handle_missing_values,
     _remove_duplicates,
-    _select_and_rename_column,
+    _rename_column,
     _validate_dataframe,
+    clean_apple_data,
 )
 
 
@@ -29,7 +30,8 @@ def raw_data():
 
 
 def test_select_and_rename_column(raw_data):
-    result = _select_and_rename_column(raw_data)
+    """Test that the 'Close' column is correctly selected and renamed 'close_price'."""
+    result = _rename_column(raw_data)
     assert list(result.columns) == ["close_price"]
     pd.testing.assert_series_equal(
         result["close_price"], raw_data["Close"], check_names=False
@@ -40,22 +42,30 @@ def test_select_and_rename_column(raw_data):
 
 
 def test_handle_missing_values(raw_data):
+    """Test that missing values are handled correctly."""
     result = _handle_missing_values(raw_data)
-    assert result.isna().sum().sum() == 0
-    assert result.loc[1, "Close"] == expected_value_missing
+    assert all(
+        [
+            result.isna().sum().sum() == 0,
+            result.loc[1, "Close"] == expected_value_missing,
+        ]
+    )
 
 
 expected_length_index = 4
 
 
 def test_remove_duplicates(raw_data):
+    """Test that duplicate indices are removed, ensuring a unique datetime index."""
     df = raw_data.copy().set_index("Date")
     result = _remove_duplicates(df)
-    assert result.index.duplicated().sum() == 0
-    assert len(result) == expected_length_index
+    assert all(
+        [result.index.duplicated().sum() == 0, len(result) == expected_length_index]
+    )
 
 
 def test_convert_to_numeric(raw_data):
+    """Test that all DataFrame columns are converted to numeric data types."""
     df = raw_data.copy().astype(str)
     result = _convert_to_numeric(df)
     for col in result.columns:
@@ -64,6 +74,7 @@ def test_convert_to_numeric(raw_data):
 
 
 def test_validate_dataframe(raw_data):
+    """Test that the DataFrame validation correctly checks for the 'Close' column."""
     _validate_dataframe(raw_data)
     df_missing = raw_data.drop(columns=["Close"])
     with pytest.raises(
@@ -73,6 +84,33 @@ def test_validate_dataframe(raw_data):
 
 
 def test_convert_to_datetime(raw_data):
+    """Test that the 'Date' column is correctly converted to a datetime index."""
     result = _convert_to_datetime(raw_data.copy())
-    assert isinstance(result.index, pd.DatetimeIndex)
-    assert result.index[0] == pd.Timestamp("2022-01-01")
+    assert all(
+        [
+            isinstance(result.index, pd.DatetimeIndex),
+            result.index[0] == pd.Timestamp("2022-01-01"),
+        ]
+    )
+
+
+def test_clean_apple_data_structure(raw_data):
+    """Test that the cleaned DataFrame maintains the correct structure."""
+    result = clean_apple_data(raw_data)
+    assert isinstance(result, pd.DataFrame)
+    assert "close_price" in result.columns
+
+
+def test_clean_apple_data_no_missing_values(raw_data):
+    """Test that the cleaned DataFrame contains no missing values."""
+    result = clean_apple_data(raw_data)
+    assert result.isna().sum().sum() == 0
+
+
+def test_clean_apple_data_validates_input():
+    """Test that a ValueError is raised if 'Close' column is missing."""
+    df_missing = pd.DataFrame({"Date": ["2022-01-01"], "Open": [21]})
+    with pytest.raises(
+        ValueError, match="The DataFrame does not contain a 'Close' column."
+    ):
+        clean_apple_data(df_missing)