import pytest from apollo_utils.core.constants import ALL from apollo_utils.core.constants.dsp import DSP from apollo_utils.core.constants.market import Market from datetime import timedelta from http import HTTPStatus from pytest_mock import MockerFixture from typing import Optional from server.client.clients.dsp_api import DspApiClient from server.client.utils import str_to_date from server.utils.common import get_last_date_by_weekday def get_country_codes(dsp: Optional[str] = None, include_all: bool = True, include_worldwide: bool = True) -> dict: country_code_list_spotify = ("us", "gb", "ca", "de", "it", "es", "mx", "br") country_code_list_apple = ("us", "gb", "ca", "de", "it", "es", "ro", "hu", "cn") country_code_list_amazon = ("us", "gb", "ca", "de", "it", "ma", "co", "eg") result = { DSP.SPOTIFY.value: country_code_list_spotify, DSP.APPLE.value: country_code_list_apple, DSP.AMAZON.value: country_code_list_amazon, } if include_all: result[ALL] = tuple( set(country_code_list_spotify) | set(country_code_list_apple) | set(country_code_list_amazon) ) if include_worldwide: for key, value in result.items(): result[key] = [Market.WORLDWIDE] + list(value) return {dsp: result[dsp]} if dsp else result def get_streams_insights( dsp: Optional[str] = None, include_worldwide: bool = True, include_zero_weeks: bool = False ) -> dict[str, dict[str, list]]: result = {} for dsp_index, (dsp, country_code_list) in enumerate( get_country_codes(dsp, include_all=True, include_worldwide=include_worldwide).items() ): result_item = [] for cc_index, country_code in enumerate(country_code_list): if not include_zero_weeks and (dsp_index + cc_index % 4) == 0: continue result_item.append( { "peak_streams_date": f"2023-04-1{(dsp_index + cc_index) % 10}", "peak_streams_number": 1020 * dsp_index + cc_index * 14, "strike_weeks": dsp_index + cc_index % 4, "country_code": country_code, } ) result["all_dsps" if dsp == ALL else dsp] = result_item return result def get_markets_ranks(as_map: bool = False) -> dict: return { dsp: ( {country_code: dsp_index + cc_index for cc_index, country_code in enumerate(country_code_list, 1)} if as_map else [ {"country_code": country_code, "rank": dsp_index + cc_index} for cc_index, country_code in enumerate(country_code_list, 1) ] ) for dsp_index, (dsp, country_code_list) in enumerate(get_country_codes(include_all=False).items()) } def get_streams(dsp: str, include_worldwide: bool = True) -> list[dict]: latest_date = get_last_date_by_weekday(3) return [ { "date": (str_to_date(latest_date) - timedelta(days=date_index)).isoformat(), "country_code": country_code, "dsp": dsp, "streams": 5020 + cc_index * 124 + date_index * 18, } for dsp, country_code_list in get_country_codes(dsp=dsp, include_worldwide=include_worldwide).items() for cc_index, country_code in enumerate(country_code_list) for date_index in range(14 - cc_index) ] def count_streams(country_code_map: dict, country_code: str, date_range: tuple[int, int]) -> int: count = 0 for country_code_list in country_code_map.values(): cc_index = country_code_list.index(country_code) if cc_index < 0: continue for date_index in range(max(cc_index, date_range[0]), date_range[1] + 1): count += 5020 + cc_index * 124 + (13 - date_index) * 18 return count def fix_rank_map(rank_map: dict[str, int], country_code_list: list[str]) -> dict[str, int]: missing_list = [i for i in country_code_list if i not in rank_map] if missing_list: index = (max(rank_map.values()) if rank_map else 0) + 1 for item in sorted(missing_list): rank_map[item] = index index += 1 return rank_map def get_expected_result(dsp: str, include_worldwide: bool = True, include_zero_weeks: bool = False) -> list[dict]: result = get_streams_insights(dsp, include_worldwide, include_zero_weeks) result = list(result.values())[0] rank_map = get_markets_ranks(as_map=True) rank_map = rank_map[DSP.SPOTIFY.value if dsp == ALL else dsp] rank_map = fix_rank_map(rank_map, [i["country_code"] for i in result]) country_code_map = get_country_codes(dsp=dsp, include_worldwide=include_worldwide) for item in result: country_code = item["country_code"] streams_last_week = count_streams(country_code_map, country_code, (7, 13)) streams_previous_week = count_streams(country_code_map, country_code, (0, 6)) item.update( { "rank": rank_map[country_code], "streams_last_day": count_streams(country_code_map, country_code, (13, 13)), "streams_last_week": streams_last_week, "streams_previous_week": streams_previous_week, "trend": ( round((streams_last_week - streams_previous_week) * 100 / streams_previous_week) if streams_previous_week else None ), } ) return result @pytest.mark.parametrize( "params,status", ( ({}, HTTPStatus.BAD_REQUEST), ({"isrc": "AABBCC123456"}, HTTPStatus.OK), ({"isrc": "AABBCC123456", "dsp": "spotify"}, HTTPStatus.OK), ({"isrc": "AABBCC123456", "dsp": "apple"}, HTTPStatus.OK), ({"isrc": "AABBCC123456", "dsp": "amazon", "include_worldwide": "true"}, HTTPStatus.OK), ({"isrc": "AABBCC123456", "dsp": "spotify", "include_zero_weeks": "true"}, HTTPStatus.OK), ), ) async def test_get_tracks_insights(params: dict, status: int, mocker: MockerFixture, auth: dict, client): dsp, include_worldwide, include_zero_weeks = ( params.get("dsp", ALL), params.get("include_worldwide", "false") == "true", params.get("include_zero_weeks", "false") == "true", ) async def dsp_send(url: str, *args, **kwargs): if url == "api/delphi/streams/insights": return get_streams_insights(dsp, include_worldwide) elif url == "api/delphi/streams/country-code/rank": return get_markets_ranks() elif url == "api/delphi/streams": return get_streams(dsp, include_worldwide) else: raise NotImplementedError() dsp_send_mock = mocker.patch.object(DspApiClient, "_send_request", side_effect=dsp_send) response = await client.get("/api/tracks/insights/", params=params, headers=auth) assert response.status == status assert dsp_send_mock.call_count == (3 if status == HTTPStatus.OK else 0) if status == HTTPStatus.OK: response = await response.json() assert response == get_expected_result(dsp, include_worldwide)