Skip to content
Snippets Groups Projects
test_functions.py 2.95 KiB
Newer Older
Bert Palm's avatar
Bert Palm committed
#! /usr/bin/env python
# -*- coding: utf-8 -*-

import pytest
import numpy as np
import pandas as pd

from saqc.flagger.baseflagger import BaseFlagger
from saqc.flagger.dmpflagger import DmpFlagger
from saqc.flagger.simpleflagger import SimpleFlagger

from saqc.funcs.functions import flagRange, flagSesonalRange, forceFlags, clearFlags

TESTFLAGGERS = [
    BaseFlagger(['NIL', 'GOOD', 'BAD']),
    DmpFlagger(),
    SimpleFlagger()]


@pytest.mark.parametrize('flagger', TESTFLAGGERS)
def test_range(flagger):
    # prepare
    field = 'testdata'
    index = pd.date_range(start='2011-01-01', end='2011-01-02', periods=100)
    data = pd.DataFrame(data={field: np.linspace(0, index.size - 1, index.size)}, index=index)
    flags = flagger.initFlags(data)
    # test
    data, flags = flagRange(data, flags, field, flagger, min=10, max=90)
    flagged = flagger.isFlagged(flags[field])
    assert len(flags[flagged]) == 10 + 10


@pytest.mark.parametrize('flagger', TESTFLAGGERS)
def test_flagSesonalRange(flagger):
    # prepare
    field = 'testdata'
    index = pd.date_range(start='2011-01-01', end='2014-12-31', freq='1d')
Bert Palm's avatar
Bert Palm committed
    d = [(x % 2) * 50 for x in range(index.size)]
    data = pd.DataFrame(data={field: d}, index=index)
Bert Palm's avatar
Bert Palm committed
    flags = flagger.initFlags(data)

    # test
Bert Palm's avatar
Bert Palm committed
    kwargs = dict(min=1, max=100, startmonth=7, startday=1, endmonth=8, endday=31)
    data, flags = flagSesonalRange(data, flags, field, flagger, **kwargs)
Bert Palm's avatar
Bert Palm committed
    flagged = flagger.isFlagged(flags[field])
Bert Palm's avatar
Bert Palm committed
    assert len(flags[flagged]) == (31 + 31) * 4 / 2
Bert Palm's avatar
Bert Palm committed

    flags = flagger.initFlags(data)
Bert Palm's avatar
Bert Palm committed
    kwargs = dict(min=1, max=100, startmonth=12, startday=16, endmonth=1, endday=15)
    _, flags = flagSesonalRange(data, flags, field, flagger, **kwargs)
Bert Palm's avatar
Bert Palm committed
    flagged = flagger.isFlagged(flags[field])
Bert Palm's avatar
Bert Palm committed
    assert len(flags[flagged]) == 31 * 4 / 2
Bert Palm's avatar
Bert Palm committed


@pytest.mark.parametrize('flagger', TESTFLAGGERS)
def test_clearFlags(flagger):
    # prepare
    field = 'testdata'
    index = pd.date_range(start='2011-01-01', end='2011-01-10', freq='1d')
    data = pd.DataFrame(data={field: np.linspace(0, index.size - 1, index.size)}, index=index)
    orig = flagger.initFlags(data)
    flags = orig.copy()
    # test
    flags = flagger.setFlags(flags, field)
    assert (orig != flags).all
    _, cleared = clearFlags(data, flags, field, flagger)
    assert (orig == cleared).all


@pytest.mark.parametrize('flagger', TESTFLAGGERS)
def test_forceFlags(flagger):
    # prepare
    field = 'testdata'
    index = pd.date_range(start='2011-01-01', end='2011-01-10', freq='1d')
    data = pd.DataFrame(data={field: np.linspace(0, index.size - 1, index.size)}, index=index)
    flags = flagger.initFlags(data)
    flags = flagger.setFlags(flags, field)
    orig = flags.copy()
    # test
    _, foreced = forceFlags(data, flags, field, flagger, flag=flagger.GOOD)
    assert (orig != foreced).all


Bert Palm's avatar
Bert Palm committed
if __name__ == '__main__':
    for f in TESTFLAGGERS:
        test_range(f)
        test_flagSesonalRange(f)
        test_clearFlags(f)
        test_forceFlags(f)