test_df_operations.py 1.23 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
from unittest import TestCase

import pandas as pd
from pandas.util.testing import assert_series_equal

from metagenedb.common.utils.df_operations import get_mask


class TestGetMask(TestCase):

    def setUp(self):
        self.df = pd.DataFrame(
            [
                ['breizh', 6],
                ['loquemo', 7],
                ['sea', 3],
            ],
            columns=['name', 'length']
        )

    def test_get_mask(self):
        range = (5, 10)
        expected_series = pd.Series(
            [True, True, False],
            name='length'
        )
        test_mask = get_mask(self.df, range, 'length')
        assert_series_equal(test_mask, expected_series)

    def test_get_mask_lower_equal(self):
        range = (3, 5)
        expected_series = pd.Series(
            [False, False, True],
            name='length'
        )
        test_mask = get_mask(self.df, range, 'length')
        assert_series_equal(test_mask, expected_series)

    def test_get_mask_upper_equal(self):
        range = (4, 6)
        expected_series = pd.Series(
            [False, False, False],
            name='length'
        )
        test_mask = get_mask(self.df, range, 'length')
        assert_series_equal(test_mask, expected_series)