"""Unit tests for demographic model layer."""

from unittest.mock import patch

import pytest

from sound_recordings.models import demographics

# This mock SQL is to test that filtering is correct
MOCK_SQL = "{filter_clause}"


@pytest.fixture
def mock_load_query():
    """Mock load query."""
    with patch("sound_recordings.models.demographics.SQLLoader") as sql_loader:
        load_query = sql_loader.load_query
        load_query.return_value = MOCK_SQL
        yield load_query


@pytest.fixture
def mock_db_response():
    """Mock db response."""
    return (
        5214236,
        18016044,
        11165474,
        5969109,
        2942418,
        1606414,
        436636,
        77677,
        8273548,
        34247706,
        2906754,
    )


@pytest.fixture
def expected_response():
    """Expect response from model."""
    return {
        "_17": 5214236,
        "_18_22": 18016044,
        "_23_27": 11165474,
        "_28_34": 5969109,
        "_35_44": 2942418,
        "_45_59": 1606414,
        "_60_": 436636,
        "UA": 77677,
        "M": 8273548,
        "F": 34247706,
        "UG": 2906754,
    }


@pytest.fixture
def permissions_filter():
    """Mock permissions filter."""
    return {"label_ids": [7123], "artist_ids": None, "subaccount_ids": None}


@pytest.fixture
def mock_snowflake_fetchone(mock_db_response):
    """Mock snowflake."""
    with patch("sound_recordings.models.demographics.snowflake.fetchone") as fetchall:
        fetchall.return_value = mock_db_response
        yield fetchall


def test_demographics_returns_results(
    mock_load_query, permissions_filter, mock_snowflake_fetchone, expected_response
):
    """Should use demographic query."""
    result = demographics.get_demographics(
        permissions_filter,
        "TEST",
        [],
        [1, 286],
        "2019-12-01",
        "2019-12-01",
        ["theorchard"],
    )
    mock_load_query.assert_called_with("demographics")
    mock_snowflake_fetchone.assert_called_once()
    assert result == expected_response


def test_countries_returns_results(
    mock_load_query, permissions_filter, mock_snowflake_fetchone, expected_response
):
    """Should use demographic country query."""
    result = demographics.get_demographics(
        permissions_filter,
        "TEST",
        ["US", "EN"],
        [1, 286],
        "2019-12-01",
        "2019-12-01",
        ["theorchard"],
    )
    mock_load_query.assert_called_with("demographics_by_country")
    mock_snowflake_fetchone.assert_called_once()
    assert result == expected_response