diff --git a/tests/data/test_projection.py b/tests/data/test_projection.py index 9133667..9b87dec 100644 --- a/tests/data/test_projection.py +++ b/tests/data/test_projection.py @@ -1,10 +1,12 @@ # Copyright (C) 2022-2023 Intel Corporation # SPDX-License-Identifier: MIT +import unittest + import pandas as pd + from p3.data import projection from p3.data._projection import _collapse -import unittest class TestProjection(unittest.TestCase): @@ -24,6 +26,18 @@ def test_collapse(self): pd.testing.assert_frame_equal(df, expected_df) + def test_collapse_null(self): + """Check that _collapse skips null values""" + data = {"c1": ["x", "y", "z"], "c2": ["1", None, "3"]} + df = pd.DataFrame(data) + + _collapse(df, ["c1", "c2"], "c3") + + expected_data = {"c3": ["x-1", "y", "z-3"]} + expected_df = pd.DataFrame(expected_data) + + pd.testing.assert_frame_equal(df, expected_df) + def test_required_columns(self): """p3.data.projection.required_columns""" df = pd.DataFrame()