import pytest

from dw_base.spark.udf.customs.indonesia_qymc_judge import is_prefix_and_concat, get_qymc_prefix


@pytest.mark.parametrize("qymc, qymc_org, expected", [
    ('PT. ABC', 'PT ABC', 'PT ABC'),
    ('PT  ABC', 'ABC', 'PT ABC'),
    ('CV. ABC', 'CV ABC', 'CV ABC'),
    ('CV ABC', 'CV ABC', 'CV ABC'),
    (None, 'ABC', 'ABC'),
    ('ABC', None, None),
    (None, None, None),
])
def test_is_prefix(qymc: str, qymc_org: str, expected: str):
    assert is_prefix_and_concat(qymc, qymc_org) == expected


@pytest.mark.parametrize("qymcs, expected", [
    (['CV ABC', 'PT ABC'], 'PT '),
    (['CV. ABC', 'ABC'], 'CV '),
    (['ABC'], None),
    (['PT. XXX','CV CV'],'PT ')
])
def test_get_qymc_prefix(qymcs: list, expected: str):
    assert get_qymc_prefix(qymcs) == expected