"""Unit tests for the set_prediction_window function."""

import copd
import pandas as pd
import pytest


@pytest.fixture
def input_df():
    """Sample input data including an exacerbation."""
    return pd.DataFrame({'Date': pd.date_range('2022-01-01', '2022-01-10'),
                         'IsExac': [0, 0, 0, 0, 0, 0, 1, 0, 0, 0]})


def test_check_correct_three_day_window_set(input_df):
    """Check the correct rows are set to exacerbations for a three day window."""
    output_df = copd.set_prediction_window(df=input_df, prediction_window=3)
    expected_df = pd.DataFrame({'Date': pd.date_range('2022-01-01', '2022-01-10'),
                                'IsExac': [0, 0, 0, 0, 1, 1, 1, 0, 0, 0]})
    pd.testing.assert_frame_equal(output_df, expected_df)


def test_check_correct_five_day_window_set(input_df):
    """Check the correct rows are set to exacerbations for a five day window."""
    output_df = copd.set_prediction_window(df=input_df, prediction_window=5)
    expected_df = pd.DataFrame({'Date': pd.date_range('2022-01-01', '2022-01-10'),
                                'IsExac': [0, 0, 1, 1, 1, 1, 1, 0, 0, 0]})
    pd.testing.assert_frame_equal(output_df, expected_df)