"""Test market share util functions."""

from feed_ingestion.util import marketshare_util


def test_cast_data():
    """Test cast_data()."""
    rows = [
        {'col1': 'aaa', 'col2': '1', 'col3': '2.3'},
        {'col1': 'bbb', 'col2': '2', 'col3': '2.33'},
    ]

    expected_rows = [
        {'col1': 'aaa', 'col2': 1, 'col3': 2.3},
        {'col1': 'bbb', 'col2': 2, 'col3': 2.33},
    ]

    res = marketshare_util.cast_data(rows, {
        'col1': str,
        'col2': int,
        'col3': float
    })

    assert res == expected_rows


def test_groupby():
    """Test groupby()."""
    rows = [
        {'country': 'UA', 'col1': 1, 'col2': 2, 'col3': 'val'},
        {'country': 'UA', 'col1': 2, 'col2': 1, 'col3': 'val'},
        {'country': 'US', 'col1': 3, 'col2': 2, 'col3': 'val'},
        {'country': 'US', 'col1': 2, 'col2': 3, 'col3': 'val'},
    ]

    expected_rows = sorted([
        {'country': 'UA', 'col1': 3, 'col2': 3},
        {'country': 'US', 'col1': 5, 'col2': 5},
    ], key=lambda x: x['country'])

    res = sorted(marketshare_util.groupby(rows, 'country', {
        'col1': sum,
        'col2': sum
    }), key=lambda x: x['country'])

    assert res == expected_rows


def test_join():
    """Test join()."""
    ds = [
        {'country': 'UA', 'col1': 1},
        {'country': 'US', 'col1': 2},
    ]

    other = [
        {'country': 'UA', 'col2': 'test'},
        {'country': 'US', 'col2': 'test1'},
    ]

    expected_rows = sorted([
        {'country': 'UA', 'col1': 1, 'col2': 'test'},
        {'country': 'US', 'col1': 2, 'col2': 'test1'},
    ], key=lambda x: x['country'])

    res = sorted(
        marketshare_util.join('country', ds, other),
        key=lambda x: x['country'])

    assert res == expected_rows