diff --git a/tests/test_forward_link.py b/tests/test_forward_link.py index 583b5485..d1c0f6cf 100644 --- a/tests/test_forward_link.py +++ b/tests/test_forward_link.py @@ -1,5 +1,7 @@ import numpy as np import pandas as pd +import pytest +from helper_functions import load_and_format from pandas.testing import assert_frame_equal, assert_series_equal from src.forward_link import calculate_imputation_link, mask_values @@ -144,195 +146,41 @@ def test_multiple_filters(self): assert_frame_equal(df_copy, expected) -class TestLink: +scenarios = ["calculate_links_test_data"] - # from scenario 33_multi_variable_C_BI_R - # We could parametrise this with more scenarios if needed - df = pd.DataFrame( - data={ - "identifier": [ - 10001, - 10001, - 10001, - 10002, - 10002, - 10002, - 10001, - 10001, - 10001, - 10002, - 10002, - 10002, - 10005, - 10005, - 10005, - ], - "date": [ - 202001, - 202002, - 202003, - 202001, - 202002, - 202003, - 202001, - 202002, - 202003, - 202001, - 202002, - 202003, - 202001, - 202002, - 202003, - ], - "group": [1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 1, 2, 2], - "question": [ - 547.0, - 362.0, - 895.0, - 381.0, - 573.0, - 214.0, - 961.0, - 267.0, - 314.0, - 555.0, - 628.0, - 736.0, - np.nan, - np.nan, - 100.0, - ], - "f_predictive_question": [ - np.nan, - 547.0, - 362.0, - np.nan, - 381.0, - 573.0, - np.nan, - 961.0, - 267.0, - np.nan, - 555.0, - 628.0, - np.nan, - np.nan, - np.nan, - ], - "b_predictive_question": [ - 362.0, - 895.0, - np.nan, - 573.0, - 214.0, - np.nan, - 267.0, - 314.0, - np.nan, - 628.0, - 736.0, - np.nan, - np.nan, - 100.0, - np.nan, - ], - "f_matched_pair": [ - False, - True, - True, - False, - True, - True, - False, - True, - True, - False, - True, - True, - False, - False, - False, - ], - "b_matched_pair": [ - True, - True, - False, - True, - True, - False, - True, - True, - False, - True, - True, - False, - False, - False, - False, - ], - } - ) - def test_forward_link(self): - - expected_f_link = pd.Series( - [ - 1.0, - 1.0075431034482758, - 1.186096256684492, - 1.0, - 1.0075431034482758, - 1.186096256684492, - 1.0, - 0.5903693931398417, - 1.1731843575418994, - 1.0, - 0.5903693931398417, - 1.1731843575418994, - 1.0, - 0.5903693931398417, - 1.1731843575418994, - ] - ) +@pytest.mark.parametrize("scenario", scenarios) +class TestLinks: + def test_forward_links(self, scenario): + """Test if function returns the f_link column""" + + df_input = load_and_format("tests/" + scenario + ".csv") - f_link = calculate_imputation_link( - self.df, - ["group", "date"], + expected_link = df_input["f_link"] + + link_to_test = calculate_imputation_link( + df_input, + ["group", "period"], "f_matched_pair", "question", "f_predictive_question", ) - assert_series_equal(f_link, expected_f_link) - - def test_backward_link(self): - - expected_b_link = pd.Series( - [ - 0.9925133689839573, - 0.8431018935978359, - 1.0, - 0.9925133689839573, - 0.8431018935978359, - 1.0, - 1.693854748603352, - 0.8523809523809524, - 1.0, - 1.693854748603352, - 0.8523809523809524, - 1.0, - 0.9925133689839573, - 0.8523809523809524, - 1.0, - ] - ) + assert_series_equal(link_to_test, expected_link, check_names=False) + + def test_back_links(self, scenario): + """Test if function returns the b_link column""" + + df_input = load_and_format("tests/" + scenario + ".csv") + + expected_link = df_input["b_link"] - b_link = calculate_imputation_link( - self.df, - ["group", "date"], + link_to_test = calculate_imputation_link( + df_input, + ["group", "period"], "b_matched_pair", "question", "b_predictive_question", ) - assert_series_equal(b_link, expected_b_link) + assert_series_equal(link_to_test, expected_link, check_names=False)