lichess-bot/test_bot/test_model.py

"""Tests for the models."""

import datetime
from lib import model
import yaml
from lib import config
from collections import defaultdict, Counter
from lib.blocklist import OnlineBlocklist
from lib.timer import Timer
from lib.lichess_types import ChallengeType, UserProfileType, GameEventType, PlayerType


def test_challenge() -> None:
    """Test the challenge model."""
    challenge: ChallengeType = {"id": "zzzzzzzz", "url": "https://lichess.org/zzzzzzzz", "status": "created",
                                "challenger": {"id": "c", "name": "c", "rating": 2000, "title": None, "online": True},
                                "destUser": {"id": "b", "name": "b", "rating": 3000, "title": "BOT", "online": True},
                                "variant": {"key": "standard", "name": "Standard", "short": "Std"}, "rated": False,
                                "speed": "bullet",
                                "timeControl": {"type": "clock", "limit": 90, "increment": 1, "show": "1.5+1"},
                                "color": "random", "finalColor": "white", "perf": {"icon": "\ue032", "name": "Bullet"}}
    user_profile: UserProfileType = {"id": "b", "username": "b",
                                     "perfs": {"bullet": {"games": 100, "rating": 3000, "rd": 150, "prog": -10},
                                               "blitz": {"games": 100, "rating": 3000, "rd": 150, "prog": -10},
                                               "rapid": {"games": 100, "rating": 3000, "rd": 150, "prog": -10},
                                               "classical": {"games": 100, "rating": 3000, "rd": 150, "prog": -10},
                                               "correspondence": {"games": 100, "rating": 3000, "rd": 150, "prog": -10},
                                               "antichess": {"games": 100, "rating": 3000, "rd": 150, "prog": -10,
                                                             "prov": True}},
                                     "title": "BOT", "createdAt": 1500000000000,
                                     "profile": {"bio": "This is my bio",
                                                 "links": "https://github.com/lichess-bot-devs/lichess-bot"},
                                     "seenAt": 1700000000000, "playTime": {"total": 1000000, "tv": 10000},
                                     "url": "https://lichess.org/@/b",
                                     "count": {"all": 600, "rated": 500, "ai": 50, "draw": 200, "drawH": 50, "loss": 50,
                                               "lossH": 50, "win": 250, "winH": 200, "bookmark": 0, "playing": 0,
                                               "import": 0, "me": 0},
                                     "followable": True, "following": False, "blocking": False}

    with open("./config.yml.default") as file:
        CONFIG = yaml.safe_load(file)
    CONFIG["token"] = ""
    CONFIG["challenge"]["allow_list"] = []
    CONFIG["challenge"]["block_list"] = []
    CONFIG["challenge"]["min_rating"] = 0
    CONFIG["challenge"]["max_rating"] = 4000
    CONFIG["challenge"]["rating_difference"] = None
    configuration = config.Configuration(CONFIG).challenge
    recent_challenges: defaultdict[str, list[Timer]] = defaultdict()
    recent_challenges["c"] = []
    online_block_list = OnlineBlocklist([])

    challenge_model = model.Challenge(challenge, user_profile)
    assert challenge_model.id == "zzzzzzzz"
    assert challenge_model.rated is False
    assert challenge_model.variant == "standard"
    assert challenge_model.speed == "bullet"
    assert challenge_model.time_control["show"] == "1.5+1"
    assert challenge_model.color == "white"
    supported = challenge_model.is_supported(configuration, recent_challenges, Counter(), online_block_list, user_profile)
    assert supported == (True, "")

    CONFIG["challenge"]["min_base"] = 120
    assert challenge_model.is_supported(configuration, recent_challenges, Counter(), online_block_list, user_profile) == (
        False,
        "timeControl",
    )


def test_challenge_rating_filters() -> None:
    """Test challenge rating filtering for incoming challenges."""
    challenge: ChallengeType = {"id": "zzzzzzzz", "url": "https://lichess.org/zzzzzzzz", "status": "created",
                                "challenger": {"id": "c", "name": "c", "rating": 2000, "title": None, "online": True},
                                "destUser": {"id": "b", "name": "b", "rating": 3000, "title": "BOT", "online": True},
                                "variant": {"key": "standard", "name": "Standard", "short": "Std"}, "rated": False,
                                "speed": "bullet",
                                "timeControl": {"type": "clock", "limit": 90, "increment": 1, "show": "1.5+1"},
                                "color": "random", "finalColor": "white", "perf": {"icon": "\ue032", "name": "Bullet"}}
    user_profile: UserProfileType = {"id": "b", "username": "b",
                                     "perfs": {"bullet": {"games": 100, "rating": 3000, "rd": 150, "prog": -10}},
                                     "title": "BOT"}

    with open("./config.yml.default") as file:
        CONFIG = yaml.safe_load(file)
    CONFIG["token"] = ""
    CONFIG["challenge"]["allow_list"] = []
    CONFIG["challenge"]["block_list"] = []
    CONFIG["challenge"]["min_rating"] = 0
    CONFIG["challenge"]["max_rating"] = 4000
    CONFIG["challenge"]["rating_difference"] = None
    configuration = config.Configuration(CONFIG).challenge
    recent_challenges: defaultdict[str, list[Timer]] = defaultdict()
    recent_challenges["c"] = []
    online_block_list = OnlineBlocklist([])

    challenge_model = model.Challenge(challenge, user_profile)

    # Default config should accept all ratings
    supported = challenge_model.is_supported(configuration, recent_challenges, Counter(), online_block_list, user_profile)
    assert supported == (True, "")

    # Test max_rating filter
    CONFIG["challenge"]["max_rating"] = 1500
    assert challenge_model.is_supported(configuration, recent_challenges, Counter(), online_block_list, user_profile) == (
        False, "generic")

    # Test min_rating filter
    CONFIG["challenge"]["max_rating"] = 4000
    CONFIG["challenge"]["min_rating"] = 2500
    assert challenge_model.is_supported(configuration, recent_challenges, Counter(), online_block_list, user_profile) == (
        False, "generic")

    # Test rating_difference filter (bot is 3000, challenger is 2000, diff is 1000)
    CONFIG["challenge"]["min_rating"] = 0
    CONFIG["challenge"]["rating_difference"] = 500
    assert challenge_model.is_supported(configuration, recent_challenges, Counter(), online_block_list, user_profile) == (
        False, "generic")

    # Rating difference large enough to accept
    CONFIG["challenge"]["rating_difference"] = 1500
    supported = challenge_model.is_supported(configuration, recent_challenges, Counter(), online_block_list, user_profile)
    assert supported == (True, "")

    # Test that rating_difference narrows the range
    # min_rating=0, max_rating=4000, but diff=500 from bot rating 3000
    CONFIG["challenge"]["rating_difference"] = 500
    CONFIG["challenge"]["min_rating"] = 0
    CONFIG["challenge"]["max_rating"] = 4000
    assert challenge_model.is_supported(configuration, recent_challenges, Counter(), online_block_list, user_profile) == (
        False, "generic")

    # Test with AI opponent (no rating) - should always accept
    CONFIG["challenge"]["rating_difference"] = None
    CONFIG["challenge"]["max_rating"] = 1000
    ai_challenge: ChallengeType = {**challenge,
                                    "challenger": {"id": "ai", "name": "AI level 5", "aiLevel": 5}}
    ai_challenge_model = model.Challenge(ai_challenge, user_profile)
    assert ai_challenge_model.is_supported_rating(configuration, user_profile) is True


def test_game() -> None:
    """Test the game model."""
    game: GameEventType = {"id": "zzzzzzzz", "variant": {"key": "standard", "name": "Standard", "short": "Std"},
                           "speed": "bullet", "perf": {"name": "Bullet"}, "rated": False, "createdAt": 1700000000000,
                           "white": {"id": "c", "name": "c", "title": None, "rating": 2000},
                           "black": {"id": "b", "name": "b", "title": "BOT", "rating": 3000},
                           "initialFen": "startpos", "clock": {"initial": 90000, "increment": 1000}, "type": "gameFull",
                           "state": {"type": "gameState", "moves": "", "wtime": 90000, "btime": 90000, "winc": 1000,
                                     "binc": 1000, "status": "started"}}
    username = "b"
    base_url = "https://lichess.org/"
    abort_time = datetime.timedelta(seconds=30)

    game_model = model.Game(game, username, base_url, abort_time)
    assert game_model.id == "zzzzzzzz"
    assert game_model.mode == "casual"
    assert game_model.is_white is False
    assert game_model.my_color == "black"
    assert game_model.url() == "https://lichess.org/zzzzzzzz/black"
    assert game_model.short_url() == "https://lichess.org/zzzzzzzz"
    assert game_model.pgn_event() == "Casual Bullet game"
    assert game_model.time_control() == "90+1"
    assert game_model.is_abortable() is True


def test_player() -> None:
    """Test the player model."""
    player: PlayerType = {"id": "b", "name": "b", "rating": 3000, "title": "BOT", "online": True}
    player_model = model.Player(player)
    assert player_model.is_bot is True
    assert str(player_model) == "BOT b (3000)"


def test_is_chess_960() -> None:
    """Test the is_chess_960 function."""
    items = [
        {"fen": "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1", "type": "standard", "is_960": False},
        {"fen": "brnkrqnb/pppppppp/8/8/8/8/PPPPPPPP/BRNKRQNB w KQkq - 0 1", "type": "960", "is_960": True},  # pos1
        {"fen": "nrbbnkqr/pppppppp/8/8/8/8/PPPPPPPP/NRBBNKQR w KQkq - 0 1", "type": "960", "is_960": True},  # pos2
    ]
    for item in items:
        fen = str(item["fen"])
        expected = item["is_960"]
        result = model.is_chess_960(fen)
        assert result == expected