poyee-data-warehouse/dw_base/spark/udf/contacts/test/ctc_common_test.py

import pytest
from dw_base.spark.udf.contacts.ctc_common import get_shh_email_status
from dw_base.spark.udf.contacts.ctc_common import clean_email_status


@pytest.mark.parametrize("inv, level, expected", [
    (True, '-8', 'low'),  # level <= -7 且 inv 为 True
    (False, '-8', 'high'),  # level <= -7 且 inv 为 False
    (True, '-1', 'low'),  # -7 < level <= 0 且 inv 为 True
    (False, '-1', 'middle'),  # -7 < level <= 0 且 inv 为 False
    (True, 'invalid', 'low'),  # 无效的 level 值（ValueError 异常）
])
def test_get_shh_email_status(inv, level, expected):
    result = get_shh_email_status(inv, level)
    assert result == expected


@pytest.mark.parametrize("source, match_level, expected", [
    ('shh', 1, 'PERFECT_MATCH'),
    ('shh', 2, 'SPECULATION_VERIFICATION'),
    ('shh', -1, 'SPECULATION_VERIFICATION'),
    ('shh', 0.95, 'POSSIBLE_MATCH'),
    ('shh', 0.5, 'LOW_MATCH'),
    ('snovio', 'valid', 'PERFECT_MATCH'),
    ('snovio', 'verified', 'PERFECT_MATCH')
])
def test_clean_email_status_functionality(source, match_level, expected):
    result = clean_email_status(source, match_level)
    assert result == expected


if __name__ == '__main__':
    pytest.main()