diff --git a/test/flagger/test_flagger.py b/test/flagger/test_flagger.py index 689d79d6e8e67f21d65d33ede6d2277598c68c8d..b3e9cb3a2a6b8d161d0662b11e0bdb94c6297490 100644 --- a/test/flagger/test_flagger.py +++ b/test/flagger/test_flagger.py @@ -8,9 +8,7 @@ import pytest import numpy as np import pandas as pd from pandas.api.types import is_bool_dtype -from pandas.core.indexing import IndexingError -from saqc.funcs.functions import flagRange, flagSesonalRange, forceFlags, clearFlags from test.common import TESTFLAGGER @@ -37,6 +35,51 @@ DATASETS = [ ] +@pytest.mark.parametrize("data", DATASETS) +@pytest.mark.parametrize("flagger", TESTFLAGGER) +def test_setFlagger(data, flagger): + + field, *_ = data.columns + + this_flagger = flagger.initFlags(data) + other_flagger = this_flagger.getFlagger(iloc=slice(None, None, 3)).setFlags(field) + result_flagger = this_flagger.setFlagger(other_flagger) + + other_flags = other_flagger.getFlags() + result_flags = result_flagger.getFlags(field) + + assert np.all( + result_flagger.getFlags(loc=other_flagger.getFlags().index) == other_flags + ) + + assert np.all( + result_flags[~result_flags.index.isin(other_flags.index)] == flagger.UNFLAGGED + ) + + +@pytest.mark.parametrize("data", DATASETS) +@pytest.mark.parametrize("flagger", TESTFLAGGER) +def test_setFlaggerDiff(data, flagger): + + field, *_ = data.columns + iloc = slice(None, None, 2) + + other_data = data.iloc[iloc] + other_data.index = other_data.index + pd.Timedelta(minutes=2, seconds=25) + + this_flagger = flagger.initFlags(data).setFlags(field, flag=flagger.BAD) + other_flagger = flagger.initFlags(other_data) + result_flagger = this_flagger.setFlagger(other_flagger) + + assert np.all( + result_flagger.getFlags(field, loc=other_data.index) + == other_flagger.getFlags(field) + ) + assert np.all( + result_flagger.getFlags(field, loc=data.index) == this_flagger.getFlags(field) + ) + + @pytest.mark.parametrize("data", DATASETS) @pytest.mark.parametrize("flagger", TESTFLAGGER) def test_initFlags(data, flagger): @@ -81,9 +124,24 @@ def test_isFlaggedDataFrame(data, flagger): df_tests = [ (flagger.isFlagged(), mask), (flagger.setFlags(field).isFlagged(), ~mask), - (flagger.setFlags(field, flag=flagger.GOOD).isFlagged(flag=flagger.GOOD, comparator=">"), mask), - (flagger.setFlags(field, flag=flagger.GOOD).isFlagged(flag=flagger.GOOD, comparator="<"), mask), - (flagger.setFlags(field, flag=flagger.GOOD).isFlagged(flag=flagger.GOOD, comparator="=="), ~mask), + ( + flagger.setFlags(field, flag=flagger.GOOD).isFlagged( + flag=flagger.GOOD, comparator=">" + ), + mask, + ), + ( + flagger.setFlags(field, flag=flagger.GOOD).isFlagged( + flag=flagger.GOOD, comparator="<" + ), + mask, + ), + ( + flagger.setFlags(field, flag=flagger.GOOD).isFlagged( + flag=flagger.GOOD, comparator="==" + ), + ~mask, + ), ] for flags, expected in df_tests: assert np.all(flags[field] == expected) @@ -106,9 +164,24 @@ def test_isFlaggedSeries(data, flagger): series_tests = [ (flagger.isFlagged(field), mask), (flagger.setFlags(field).isFlagged(field), ~mask), - (flagger.setFlags(field, flag=flagger.GOOD).isFlagged(field, flag=flagger.GOOD, comparator=">"), mask), - (flagger.setFlags(field, flag=flagger.GOOD).isFlagged(field, flag=flagger.GOOD, comparator="<"), mask), - (flagger.setFlags(field, flag=flagger.GOOD).isFlagged(field, flag=flagger.GOOD, comparator="=="), ~mask), + ( + flagger.setFlags(field, flag=flagger.GOOD).isFlagged( + field, flag=flagger.GOOD, comparator=">" + ), + mask, + ), + ( + flagger.setFlags(field, flag=flagger.GOOD).isFlagged( + field, flag=flagger.GOOD, comparator="<" + ), + mask, + ), + ( + flagger.setFlags(field, flag=flagger.GOOD).isFlagged( + field, flag=flagger.GOOD, comparator="==" + ), + ~mask, + ), ] for flags, expected in series_tests: assert np.all(flags == expected)