|
| 1 | +"""Test objects in dve.reporting.utility""" |
| 2 | +# pylint: disable=missing-function-docstring |
| 3 | + |
| 4 | +import polars as pl |
| 5 | + |
| 6 | +from dve.core_engine.backends.utilities import pl_row_count |
| 7 | +from dve.reporting.utils import extract_and_pivot_keys |
| 8 | + |
| 9 | + |
| 10 | +def test_extract_and_pivot_keys(): |
| 11 | + df = pl.DataFrame({ |
| 12 | + "entity": ["test1", "test2", "test3", "test4"], |
| 13 | + "FailureType": ["submission1", "submission2", "submission3", "submission4"], |
| 14 | + "id": [ |
| 15 | + "Key1: Value1 -- Key2: Value2 -- Key3: Value3", |
| 16 | + "Key1: Value1 -- Key2: Value2", |
| 17 | + "", |
| 18 | + None, |
| 19 | + ] |
| 20 | + }) |
| 21 | + result_df = extract_and_pivot_keys(df, key_field="id") |
| 22 | + expected_df = pl.DataFrame({ |
| 23 | + "entity": ["test1", "test2", "test3", "test4"], |
| 24 | + "FailureType": ["submission1", "submission2", "submission3", "submission4"], |
| 25 | + "Key1_Identifier": ["Value1", "Value1", None, None], |
| 26 | + "Key2_Identifier": ["Value2", "Value2", None, None], |
| 27 | + "Key3_Identifier": ["Value3", None, None, None], |
| 28 | + }) |
| 29 | + |
| 30 | + assert pl_row_count(result_df) == pl_row_count(df) |
| 31 | + assert result_df.equals(expected_df) |
| 32 | + |
| 33 | + |
| 34 | +def test_extract_and_pivot_keys_with_empty_key_field(): |
| 35 | + df = pl.DataFrame({ |
| 36 | + "entity": ["test1", "test2", "test3"], |
| 37 | + "FailureType": ["submission1", "submission2", "submission3"], |
| 38 | + "Key": ["", "None", None] |
| 39 | + }) |
| 40 | + result_df = extract_and_pivot_keys(df) |
| 41 | + |
| 42 | + assert pl_row_count(result_df) == pl_row_count(df) |
| 43 | + assert result_df.equals(df) |
0 commit comments