import pytest

from dw_base.spark.udf.customs.company_abbr import company_abbr


@pytest.mark.parametrize("country_name, company_name, expected", [
    ('germany', 'Wilhelm Manz GmbH & Co. KG', 'WILHELM MANZ'),
    ('germany', 'Wilhelm Zuleeg GmbH', 'WILHELM ZULEEG'),
    ('germany', 'Paragon GmbH & Co. KGaA', 'PARAGON GMBH AND CO KGAA'),
])
def test_company_abbr(country_name: str, company_name: str, expected: str):
    result = company_abbr(country_name, company_name)
    assert result == expected