from itertools import count from unittest.mock import AsyncMock, Mock import numpy as np import pandas as pd import pytest from pandas import DataFrame, Series from src.backend.constants import Flags, FlagTableColumns, OutputColumns from src.backend.constants import SnowFlakeColumns as SFCols from src.backend.logic.audit import flags from src.backend.logic.audit.logic import common DETAILS = FlagTableColumns.DETAILS class TestFlagChecksAT: _class = flags.FlagChecksAT @pytest.fixture def mock_at(self): """Return a mock AT DataFrame.""" at_row = { SFCols.UPC: "123", SFCols.ISRC: "456", } return DataFrame([at_row]) @pytest.mark.parametrize("has_asset_id", [True, False]) @pytest.mark.parametrize( "topic_channel,should_flag", [ ("", True), (" ", True), (" topic channel ", False), ("topic channel", False), (None, True), (np.nan, True), (float("nan"), True), ], ) def test_flag_no_topic_channel( self, mock_flagger, mocker_add_flag_db, mock_at, has_asset_id, topic_channel, should_flag, ): mock_at[SFCols.CHANNEL_DISPLAY_NAME] = topic_channel mock_at[SFCols.ASSET_ID] = "123" if has_asset_id else np.nan instance = self._class(mock_at, mock_flagger) instance.no_topic_channel() was_flagged = mock_at.loc[0][OutputColumns.AUDIT_FLAG_NO_TOPIC_CHANNEL] should_flag = should_flag or not has_asset_id assert was_flagged == should_flag if should_flag: arg0, arg1, arg2 = mocker_add_flag_db.call_args.args assert arg0.empty != should_flag assert arg1 == mock_flagger assert arg2 == Flags.NO_TOPIC_CHANNEL else: assert not mocker_add_flag_db.called @pytest.mark.parametrize("suffixed_channel_topic_name", [True, False]) @pytest.mark.parametrize( "channel_display_name,artist_name,should_flag", [ # Some real examples of various labels. None of these should be flagged! *[ (*strings, False) for strings in [ ("kordz feat. Frederic Robinson", "kordz"), ("kordz feat. Frederic Robinson", "Kordz "), ("kordz feat. Frederic Robinson", "frederic robinson "), ("Charlie Cardona", "Charlie"), ("Charlie Cardona", "charlie"), ("Charlie Cardona", " charlie "), ("Charlie Cardona", " charlie "), ("Moglii & NOVAA", "Moglii, NOVAA"), ("algoma furnace & ghost orchard", "Algoma Furnace, ghost orchard"), ("Casper the Ghost", "Casper, The Ghost"), ("xander.", "Xander"), ("Lo...", "LO"), ("Noah Slee", "Noah Slee, Shiloh Dynasty"), ("Midnight", "midnight, ylxr "), ] ], ], ) def test_flag_bad_topic_channel( self, mock_flagger, mocker_add_flag_db, mock_at, channel_display_name, artist_name, should_flag, suffixed_channel_topic_name, ): if suffixed_channel_topic_name: channel_display_name += flags.at.YTMisc.TOPIC_CHANNEL_SUFFIX mock_at[SFCols.CHANNEL_DISPLAY_NAME] = channel_display_name mock_at[SFCols.ARTIST] = artist_name instance = self._class(mock_at, mock_flagger) instance.bad_topic_channel() was_flagged = mock_at.loc[0][OutputColumns.AUDIT_FLAG_BAD_TOPIC_CHANNEL] assert was_flagged == should_flag if should_flag: arg0, arg1, arg2 = mocker_add_flag_db.call_args.args assert arg0.empty != should_flag assert arg1 == mock_flagger assert arg2 == Flags.BAD_TOPIC_CHANNEL else: assert not mocker_add_flag_db.called _flag_ownership_incomplete_mandatory_columns = [ SFCols.ABBRIVATION, SFCols.VTR_COUNTRIES, SFCols.RTR_COUNTRIES, SFCols.STR_COUNTRIES, ] def fill_flag_ownership_incomplete_mandatory_columns(self, mock_at: DataFrame): for col in self._flag_ownership_incomplete_mandatory_columns: if col not in mock_at.columns: mock_at[col] = None @pytest.mark.asyncio @pytest.mark.parametrize("has_asset_id", [True, False]) @pytest.mark.parametrize( "ownership,missing_territories,should_flag", [ ({"US"}, ["ES"], True), ({"US"}, [], False), ], ) async def test_flag_ownership_incomplete_countries( self, mocker, mock_flagger, mocker_add_flag_db, mock_at, ownership, missing_territories, has_asset_id, should_flag, ): self.fill_flag_ownership_incomplete_mandatory_columns(mock_at) mocked_ownership_incomplete_class = mocker.patch( "src.backend.logic.audit.flags.at._OwnershipIncomplete" ) mocked_ownership_incomplete_class.return_value.run = AsyncMock() mocked_ownership_incomplete_class.return_value.missing_territories = { 0: missing_territories } mock_at[SFCols.OWNERSHIP] = [ownership] mock_at[SFCols.ASSET_ID] = ["123" if has_asset_id else np.nan] mock_at[OutputColumns.AUDIT_FLAG_OWNERSHIP_INCOMPLETE] = ( should_flag if has_asset_id else pd.NA ) it_should_flag = should_flag if has_asset_id else False instance = self._class(mock_at, mock_flagger) await instance.ownership_incomplete() value = mock_at.loc[0, OutputColumns.AUDIT_FLAG_OWNERSHIP_INCOMPLETE] was_flagged = pd.notna(value) and value assert was_flagged == it_should_flag if it_should_flag: arg0, arg1, arg2 = mocker_add_flag_db.call_args.args assert arg0.empty != it_should_flag assert arg1 == mock_flagger assert arg2 == Flags.OWNERSHIP_INCOMPLETE assert mocker_add_flag_db.call_args.kwargs[DETAILS] else: assert not mocker_add_flag_db.called class TestOwnershipIncomplete: _class = flags.at._OwnershipIncomplete @pytest.fixture def instance(self): df = DataFrame( [ {SFCols.ASSET_ID: "123", SFCols.OWNERSHIP: "ownership1"}, ] ) return self._class(df) @pytest.fixture def patcher_is_ownership_valid(self, mocker): return mocker.patch.object( common, "is_ownership_valid", autospec=True, ) def test_ignore_territories_in_cols_cant_be_empty(self): assert len(self._class._ignore_territories_in_cols) > 0 @pytest.mark.asyncio async def test_run(self, instance, mocker): mocker.patch.object( common, "is_ownership_valid", return_value=AsyncMock(return_value=True), ) await instance.run() assert ( instance.at.loc[0, OutputColumns.AUDIT_FLAG_OWNERSHIP_INCOMPLETE] == False # noqa: E712 ) @pytest.mark.asyncio @pytest.mark.parametrize("has_asset_id", [True, False]) async def test_run(self, mocker, instance, has_asset_id): instance.at.loc[0, SFCols.ASSET_ID] = "123" if has_asset_id else np.nan mocker_handle_row = mocker.patch.object(instance, "_handle_row", autospec=True) await instance.run() assert mocker_handle_row.called == has_asset_id @pytest.mark.asyncio @pytest.mark.parametrize("ownership_valid", [True, False]) async def test_handle_row_ownership( self, mocker, patcher_is_ownership_valid, instance, ownership_valid ): idx = 0 instance.at.loc[idx, SFCols.OWNERSHIP] = {"US"} mocker.patch.object(instance, "_get_row_territories_to_ignore", autospec=True) mock_missing_territories = Mock() patcher_is_ownership_valid.return_value = ( ownership_valid, mock_missing_territories, ) await instance._handle_row(idx) assert instance.at.loc[idx, instance._flag_col] is not ownership_valid assert patcher_is_ownership_valid.called is_flagged = instance.at.loc[idx, instance._flag_col] is True assert is_flagged is not ownership_valid assert instance.missing_territories.get(idx) == ( None if ownership_valid else mock_missing_territories ), "Missing territories should be persisted if ownership is invalid" @pytest.mark.asyncio async def test_handle_row_ownership_not_string( self, patcher_is_ownership_valid, instance ): idx = 0 instance.at.loc[idx, SFCols.OWNERSHIP] = None await instance._handle_row(idx) assert instance.at.loc[idx, instance._flag_col] is True assert not patcher_is_ownership_valid.called def test_flag_row(self, instance): idx = 0 assert instance.at.loc[idx, instance._flag_col] is not True instance._flag_row(idx) assert instance.at.loc[idx, instance._flag_col] is True @pytest.mark.parametrize("test_nans", [True, False]) def test_get_row_territories_to_ignore(self, instance, test_nans): assert len(instance._ignore_territories_in_cols) > 0 numbers = count() instance.ignore_territories = {str(next(numbers)) for _ in range(3)} mock_row = Series( { col: pd.NA if test_nans else set(str(next(numbers)) for _ in range(3)) for col in instance._ignore_territories_in_cols } ) result = instance._get_row_territories_to_ignore(mock_row) expected = set(str(number) for number in range(0, next(numbers))) assert result == expected assert isinstance(result, set)