diff --git a/tests/aircall/test_aircall.py b/tests/aircall/test_aircall.py index 5f872e3a1..a49d15830 100644 --- a/tests/aircall/test_aircall.py +++ b/tests/aircall/test_aircall.py @@ -172,7 +172,7 @@ def test__retrieve_data_no_teams_case(con, mocker): # must have calls and still have a team column even if everything is NaN assert run_fetches_mock.call_count == 1 assert df.shape == (10, 11) - assert df['team'].isna().any() + assert df['team'].isin(['NO TEAM']).all() def test__retrieve_tags_from_fetch(con, mocker): diff --git a/tests/aircall/test_helpers.py b/tests/aircall/test_helpers.py index 322576c48..275a9652e 100644 --- a/tests/aircall/test_helpers.py +++ b/tests/aircall/test_helpers.py @@ -119,7 +119,7 @@ def test_build_calls_df(): fake_list_of_data_3 = [empty_df, empty_var_df, calls_df] df_3 = build_df('calls', fake_list_of_data_3) assert df_3.shape == (10, 11) - assert df_3['team'].isna().all() + assert df_3['team'].isin(['NO TEAM']).all() # empty arrays fake_list_of_data_4 = [empty_df, empty_var_df, empty_var_df] diff --git a/toucan_connectors/aircall/helpers.py b/toucan_connectors/aircall/helpers.py index 9e53a2d99..4dacaedcf 100644 --- a/toucan_connectors/aircall/helpers.py +++ b/toucan_connectors/aircall/helpers.py @@ -64,6 +64,7 @@ def build_df(dataset: str, list_of_data: List[dict]) -> pd.DataFrame: answered_at=lambda t: pd.to_datetime(t['answered_at'], unit='s'), ended_at=lambda t: pd.to_datetime(t['ended_at'], unit='s'), day=lambda t: t['ended_at'].astype(str).str[:10], + team=lambda x: x['team'].replace({np.NaN: 'NO TEAM'}), ) return total_df[COLUMN_DICTIONARY[dataset]]