71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
import numpy as np
|
|
|
|
from template_code_location.dataframe_level_anonymisation.utils import (
|
|
parse_value_list,
|
|
normalize_hierarchy_levels,
|
|
)
|
|
|
|
|
|
# ------------------------------------
|
|
# Tests for parse_value_list
|
|
# ------------------------------------
|
|
def test_parse_value_list_all_strings_digits():
|
|
values = ["1", "2", "3"]
|
|
assert parse_value_list(values) == [1, 2, 3]
|
|
|
|
|
|
def test_parse_value_list_mixed_values():
|
|
values = ["1", 2, "abc", "5"]
|
|
assert parse_value_list(values) == [1, 2, "abc", 5]
|
|
|
|
|
|
def test_parse_value_list_no_digits():
|
|
values = ["a", "b", "c"]
|
|
assert parse_value_list(values) == ["a", "b", "c"]
|
|
|
|
|
|
# ------------------------------------
|
|
# Tests for normalize_hierarchy_levels
|
|
# ------------------------------------
|
|
def test_normalize_hierarchy_levels_level_0_converted_to_numpy_array():
|
|
hierarchy = {"age": {"0": ["1", "2", "3"], "1": ["0-10", "11-20"]}}
|
|
|
|
normalized = normalize_hierarchy_levels(hierarchy)
|
|
|
|
assert "age" in normalized
|
|
assert 0 in normalized["age"]
|
|
assert isinstance(normalized["age"][0], np.ndarray)
|
|
assert normalized["age"][0].tolist() == [1, 2, 3] # converted via parse_value_list
|
|
assert normalized["age"][1] == ["0-10", "11-20"] # untouched
|
|
|
|
|
|
def test_normalize_hierarchy_levels_multiple_columns():
|
|
hierarchy = {"age": {"0": ["10", "20"]}, "gender": {"0": ["M", "F"], "1": ["*"]}}
|
|
|
|
normalized = normalize_hierarchy_levels(hierarchy)
|
|
|
|
# First column
|
|
assert isinstance(normalized["age"][0], np.ndarray)
|
|
assert normalized["age"][0].tolist() == [10, 20]
|
|
|
|
# Second column
|
|
assert isinstance(normalized["gender"][0], np.ndarray)
|
|
assert normalized["gender"][0].tolist() == ["M", "F"]
|
|
assert normalized["gender"][1] == ["*"]
|
|
|
|
|
|
def test_normalize_hierarchy_levels_mixed_digit_non_digit_at_level_0():
|
|
hierarchy = {"test": {"0": ["1", "x", "3"]}}
|
|
|
|
normalized = normalize_hierarchy_levels(hierarchy)
|
|
|
|
assert isinstance(normalized["test"][0], np.ndarray)
|
|
assert normalized["test"][0].tolist() == ["1", "x", "3"]
|
|
|
|
|
|
def test_normalize_hierarchy_levels_empty_mapping():
|
|
hierarchy = {"col": {}}
|
|
normalized = normalize_hierarchy_levels(hierarchy)
|
|
|
|
assert normalized == {"col": {}}
|