from flows.spotify import config
from flows.spotify import tasks


def test_split_countries():
    groups = ['A', 'E', 'I', 'O', ]
    to_remove = set(config.countries) - set(config.expected_countries)
    optional_countries = [v for v in config.countries + ['11'] if v not in to_remove]

    result = tasks.split_by_prefixes(prefixes=groups, data=optional_countries)

    counts = {'A': 19, 'E': 11, 'I': 15, 'O': 14}  # config.exepected_countries
    counts = {'A': 65, 'E': 38, 'I': 71, 'O': 80}  # config.countries
    counts = {'A': 19, 'E': 11, 'I': 15, 'O': 14}  # optional_countries
    post_prcessed = {k: len(v) for k,v in result.items()}
    assert post_prcessed == counts