data-validation-engine/tests/test_error_reporting/test_utils.py at 9d5e5c3c5f150f65dd643158415930bd991d35cf · NHSDigital/data-validation-engine · GitHub

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
"""Test objects in dve.reporting.utility"""
# pylint: disable=missing-function-docstring

import polars as pl

from dve.core_engine.backends.utilities import pl_row_count
from dve.reporting.utils import extract_and_pivot_keys


def test_extract_and_pivot_keys():
    df = pl.DataFrame({
        "entity": ["test1", "test2", "test3", "test4"],
        "FailureType": ["submission1", "submission2", "submission3", "submission4"],
        "id": [
            "Key1: Value1 -- Key2: Value2 -- Key3: Value3",
            "Key1: Value1 -- Key2: Value2",
            "",
            None,
        ]
    })
    result_df = extract_and_pivot_keys(df, key_field="id")
    expected_df = pl.DataFrame({
        "entity": ["test1", "test2", "test3", "test4"],
        "FailureType": ["submission1", "submission2", "submission3", "submission4"],
        "Key1_Identifier": ["Value1", "Value1", None, None],
        "Key2_Identifier": ["Value2", "Value2", None, None],
        "Key3_Identifier": ["Value3", None, None, None],
    })

    assert pl_row_count(result_df) == pl_row_count(df)
    assert result_df.equals(expected_df)


def test_extract_and_pivot_keys_with_empty_key_field():
    df = pl.DataFrame({
        "entity": ["test1", "test2", "test3"],
        "FailureType": ["submission1", "submission2", "submission3"],
        "Key": ["", "None", None]
    })
    result_df = extract_and_pivot_keys(df)

    assert pl_row_count(result_df) == pl_row_count(df)
    assert result_df.equals(df)