training/tests/test_set_pro_exac_dates.py

"""Unit tests for the set_pro_exac_dates function."""

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


@pytest.fixture
def input_df():
    """Sample input data.

    Input data covers the following cases:
    1. Duplicate non-exacerbation response (input rows 0 and 3)
    2. Exacerbation with known date (row 1)
    3. Exacerbation with unknown date (row 2)
    """
    return pd.DataFrame({'PatientId': [1, 2, 3, 1],
                         'SymptomDiaryQ11a': [1, 2, np.nan, 1],
                         'SymptomDiaryQ11b': [np.nan, pd.to_datetime('2022-01-01'),
                                              np.nan, np.nan],
                         'SubmissionTime': pd.to_datetime(['2022-01-03', '2022-01-05',
                                                           '2022-01-06', '2022-01-03']),
                         'IsCommExac': [1, 1, 0, 1]})


@pytest.fixture
def expected_df():
    """Define expected output dataframe."""
    return pd.DataFrame({'PatientId': [2, 3, 1],
                         'SymptomDiaryQ11a': [2, np.nan, 1],
                         'SymptomDiaryQ11b': [pd.to_datetime('2022-01-01'), np.nan,
                                              np.nan],
                         'SubmissionTime': pd.to_datetime(['2022-01-05', '2022-01-06',
                                                          '2022-01-03']),
                         'IsCommExac': [1, 0, 1],
                         'DateOfEvent': pd.to_datetime(['2022-01-01', '2022-01-06',
                                                        '2022-01-03'], utc=True
                                                       ).normalize(),
                         'ExacDateUnknown': [0, 0, 1]})


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