from __future__ import annotations import os import time import zoneinfo from collections.abc import Callable from datetime import UTC, datetime from decimal import Decimal from uuid import uuid4 import pytest from generalresearch.managers.leaderboard.manager import LeaderboardManager from generalresearch.managers.leaderboard.tasks import hit_leaderboards from generalresearch.models.thl.definitions import Status from generalresearch.models.thl.leaderboard import ( LeaderboardCode, LeaderboardFrequency, LeaderboardRow, ) from generalresearch.models.thl.product import ( PayoutConfig, PayoutTransformation, PayoutTransformationPercentArgs, Product, ) from generalresearch.models.thl.session import Session from generalresearch.models.thl.user import User from generalresearch.redis_helper import RedisConfig # random uuid for leaderboard tests product_id = uuid4().hex @pytest.fixture(autouse=True) def set_timezone(): os.environ["TZ"] = "UTC" time.tzset() yield # Optionally reset to default os.environ.pop("TZ", None) time.tzset() @pytest.fixture def session_factory(): return _create_session def _create_session( product_user_id: str = "aaa", country_iso: str = "us", user_payout: Decimal = Decimal("1.00"), ): user = User( product_id=product_id, product_user_id=product_user_id, ) user.product = Product( id=product_id, name="test", redirect_url="https://www.example.com", payout_config=PayoutConfig( payout_transformation=PayoutTransformation( f="payout_transformation_percent", kwargs=PayoutTransformationPercentArgs(pct=0.5), ) ), ) session = Session( user=user, started=datetime(2025, 2, 5, 6, tzinfo=UTC), id=1, country_iso=country_iso, status=Status.COMPLETE, payout=Decimal("2.00"), user_payout=user_payout, ) return session @pytest.fixture(scope="function") def setup_leaderboards(thl_redis: RedisConfig) -> Callable[..., None]: def _inner(): complete_count = { "aaa": 10, "bbb": 6, "ccc": 6, "ddd": 6, "eee": 2, "fff": 1, "ggg": 1, } sum_payout = {"aaa": 345, "bbb": 100, "ccc": 100} max_payout = sum_payout country_iso = "us" for freq in [ LeaderboardFrequency.DAILY, LeaderboardFrequency.WEEKLY, LeaderboardFrequency.MONTHLY, ]: m = LeaderboardManager( redis_client=thl_redis, board_code=LeaderboardCode.COMPLETE_COUNT, freq=freq, product_id=product_id, country_iso=country_iso, within_time=datetime(2025, 2, 5, 12, 12, 12, tzinfo=UTC), ) thl_redis.delete(m.key) thl_redis.zadd(m.key, complete_count) m = LeaderboardManager( redis_client=thl_redis, board_code=LeaderboardCode.SUM_PAYOUTS, freq=freq, product_id=product_id, country_iso=country_iso, within_time=datetime(2025, 2, 5, 12, 12, 12, tzinfo=UTC), ) thl_redis.delete(m.key) thl_redis.zadd(m.key, sum_payout) m = LeaderboardManager( redis_client=thl_redis, board_code=LeaderboardCode.LARGEST_PAYOUT, freq=freq, product_id=product_id, country_iso=country_iso, within_time=datetime(2025, 2, 5, 12, 12, 12, tzinfo=UTC), ) thl_redis.delete(m.key) thl_redis.zadd(m.key, max_payout) return _inner class TestLeaderboards: def test_leaderboard_manager( self, setup_leaderboards: Callable[..., None], thl_redis: RedisConfig ): setup_leaderboards() country_iso = "us" board_code = LeaderboardCode.COMPLETE_COUNT freq = LeaderboardFrequency.DAILY m = LeaderboardManager( redis_client=thl_redis, board_code=board_code, freq=freq, product_id=product_id, country_iso=country_iso, within_time=datetime(2025, 2, 5, 0, 0, 0, tzinfo=UTC), ) lb = m.get_leaderboard() assert lb.period_start_local == datetime( 2025, 2, 5, 0, 0, 0, tzinfo=zoneinfo.ZoneInfo(key="America/New_York") ) assert lb.period_end_local == datetime( 2025, 2, 5, 23, 59, 59, 999999, tzinfo=zoneinfo.ZoneInfo(key="America/New_York"), ) assert lb.period_start_utc == datetime(2025, 2, 5, 5, tzinfo=UTC) assert lb.row_count == 7 assert lb.rows == [ LeaderboardRow(bpuid="aaa", rank=1, value=10), LeaderboardRow(bpuid="bbb", rank=2, value=6), LeaderboardRow(bpuid="ccc", rank=2, value=6), LeaderboardRow(bpuid="ddd", rank=2, value=6), LeaderboardRow(bpuid="eee", rank=5, value=2), LeaderboardRow(bpuid="fff", rank=6, value=1), LeaderboardRow(bpuid="ggg", rank=6, value=1), ] def test_leaderboard_manager_bpuid( self, setup_leaderboards: Callable[..., None], thl_redis: RedisConfig ): setup_leaderboards() country_iso = "us" board_code = LeaderboardCode.COMPLETE_COUNT freq = LeaderboardFrequency.DAILY m = LeaderboardManager( redis_client=thl_redis, board_code=board_code, freq=freq, product_id=product_id, country_iso=country_iso, within_time=datetime(2025, 2, 5, 12, 12, 12, tzinfo=UTC), ) lb = m.get_leaderboard(bp_user_id="fff", limit=1) # TODO: this won't work correctly if I request bpuid 'ggg', because it # is ordered at the end even though it is tied, so it won't get a # row after ('fff') assert lb.rows == [ LeaderboardRow(bpuid="eee", rank=5, value=2), LeaderboardRow(bpuid="fff", rank=6, value=1), LeaderboardRow(bpuid="ggg", rank=6, value=1), ] lb.censor() assert lb.rows[0].bpuid == "ee*" def test_leaderboard_hit( self, setup_leaderboards: Callable[..., None], session_factory: Callable[..., Session], thl_redis: RedisConfig, ): setup_leaderboards() hit_leaderboards(redis_client=thl_redis, session=session_factory()) for freq in [ LeaderboardFrequency.DAILY, LeaderboardFrequency.WEEKLY, LeaderboardFrequency.MONTHLY, ]: m = LeaderboardManager( redis_client=thl_redis, board_code=LeaderboardCode.COMPLETE_COUNT, freq=freq, product_id=product_id, country_iso="us", within_time=datetime(2025, 2, 5, 12, 12, 12, tzinfo=UTC), ) lb = m.get_leaderboard(limit=1) assert lb.row_count == 7 assert lb.rows == [LeaderboardRow(bpuid="aaa", rank=1, value=11)] m = LeaderboardManager( redis_client=thl_redis, board_code=LeaderboardCode.LARGEST_PAYOUT, freq=freq, product_id=product_id, country_iso="us", within_time=datetime(2025, 2, 5, 12, 12, 12, tzinfo=UTC), ) lb = m.get_leaderboard(limit=1) assert lb.row_count == 3 assert lb.rows == [LeaderboardRow(bpuid="aaa", rank=1, value=345)] m = LeaderboardManager( redis_client=thl_redis, board_code=LeaderboardCode.SUM_PAYOUTS, freq=freq, product_id=product_id, country_iso="us", within_time=datetime(2025, 2, 5, 12, 12, 12, tzinfo=UTC), ) lb = m.get_leaderboard(limit=1) assert lb.row_count == 3 assert lb.rows == [LeaderboardRow(bpuid="aaa", rank=1, value=345 + 100)] def test_leaderboard_hit_new_row( self, setup_leaderboards: Callable[..., None], session_factory: Callable[..., None], thl_redis: RedisConfig, ): setup_leaderboards() session = session_factory(product_user_id="zzz") hit_leaderboards(redis_client=thl_redis, session=session) m = LeaderboardManager( redis_client=thl_redis, board_code=LeaderboardCode.COMPLETE_COUNT, freq=LeaderboardFrequency.DAILY, product_id=product_id, country_iso="us", within_time=datetime(2025, 2, 5, 12, 12, 12, tzinfo=UTC), ) lb = m.get_leaderboard() assert lb.row_count == 8 assert LeaderboardRow(bpuid="zzz", value=1, rank=6) in lb.rows def test_leaderboard_country(self, thl_redis: RedisConfig): m = LeaderboardManager( redis_client=thl_redis, board_code=LeaderboardCode.COMPLETE_COUNT, freq=LeaderboardFrequency.DAILY, product_id=product_id, country_iso="jp", within_time=datetime(2025, 2, 1, tzinfo=UTC), ) lb = m.get_leaderboard() assert lb.row_count == 0 assert lb.period_start_local == datetime( 2025, 2, 1, 0, 0, 0, tzinfo=zoneinfo.ZoneInfo(key="Asia/Tokyo") ) assert lb.local_start_time == "2025-02-01T00:00:00+09:00" assert lb.local_end_time == "2025-02-01T23:59:59.999999+09:00" assert lb.period_start_utc == datetime(2025, 1, 31, 15, tzinfo=UTC) print(lb.model_dump(mode="json"))