training/tests/test_remove_data_between_exacerbations.py

"""Unit tests for the remove_data_between_exacerbations function."""

import copd
import numpy as np
import pandas as pd
import pytest


@pytest.fixture
def input_df():
    """Sample input data including an exacerbation flagged for removal."""
    return pd.DataFrame({'Date': pd.date_range('2022-01-01', '2022-01-10'),
                         'IsExac': [0, 1, 0, 0, 0, 0, 1, 0, 1, 0],
                         'DaysSinceLastExac': [-1, -1, 1, 2, 3, 4, 5, 1, 2, 1],
                         'RemoveExac': [0, 0, 0, 0, 0, 0, 1, 0, 0, 0]})


@pytest.fixture
def expected_df():
    """Define expected output dataframe."""
    return pd.DataFrame({'Date': pd.date_range('2022-01-01', '2022-01-10'),
                         'IsExac': [0, 1, 0, 0, 0, 0, 1, 0, 1, 0],
                         'DaysSinceLastExac': [-1, -1, 1, 2, 3, 4, 5, 1, 2, 1],
                         'RemoveExac': [0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
                         'RemoveRow': [np.nan, np.nan, 1, 1, 1, 1, 1, np.nan,
                                       np.nan, np.nan]})


def test_output_equals_expected(input_df, expected_df):
    """Test output is as expected."""
    output_df = copd.remove_data_between_exacerbations(input_df)
    pd.testing.assert_frame_equal(output_df, expected_df)