tutor/tests/test_search.py

import unittest
from operator import or_
from functools import reduce

from hypothesis import given
from hypothesis import strategies as s

import tutor.search
from tutor.models import Color

single_colors: dict[str, set[Color]] = {
    **{c: {Color.White} for c in ["w", "white"]},
    **{c: {Color.Black} for c in ["b", "black"]},
    **{c: {Color.Blue} for c in ["u", "blue"]},
    **{c: {Color.Red} for c in ["r", "red"]},
    **{c: {Color.Green} for c in ["g", "green"]},
    **{c: {Color.Colorless} for c in ["c", "colorless"]},
}

guilds: dict[str, set[Color]] = {
    "boros": {Color.Red, Color.White},
    "golgari": {Color.Green, Color.Black},
    "selesnya": {Color.Green, Color.White},
    "dimir": {Color.Blue, Color.Black},
    "orzhov": {Color.White, Color.Black},
    "izzet": {Color.Red, Color.Blue},
    "gruul": {Color.Red, Color.Green},
    "azorius": {Color.Blue, Color.White},
    "rakdos": {Color.Black, Color.Red},
    "simic": {Color.Green, Color.Blue},
}
shards: dict[str, set[Color]] = {
    "bant": {Color.White, Color.Green, Color.Blue},
    "esper": {Color.Blue, Color.White, Color.Black},
    "grixis": {Color.Black, Color.Blue, Color.Red},
    "jund": {Color.Red, Color.Black, Color.Green},
    "naya": {Color.Green, Color.Red, Color.White},
}

all_colors: dict[str, set[Color]] = {**single_colors, **guilds, **shards}


@s.composite
def multicolor(draw: s.DrawFn) -> tuple[str, set[Color]]:
    multi = draw(s.lists(s.sampled_from("wuburgc")))
    cs: set[Color] = reduce(or_, [single_colors[c] for c in multi], set())
    text = "".join([c for c in multi])
    return text, cs


operators = s.from_type(tutor.search.Operator)


@s.composite
def colors(draw: s.DrawFn) -> tuple[str, tutor.search.Criterion]:
    key = draw(s.sampled_from(["c", "color"]))
    op = draw(operators)
    c, cs = draw(s.one_of([s.sampled_from(list(all_colors.items())), multicolor()]))
    query = f"{key}{op.value}{c}"
    return query, tutor.search.Color(op, cs)


criterion = s.one_of([colors()])


@s.composite
def searches(draw: s.DrawFn) -> tuple[str, tutor.search.Search]:
    criteria = draw(s.lists(s.one_of([colors()])))
    return " ".join([c[0] for c in criteria]), tutor.search.Search(
        [c[1] for c in criteria]
    )


class ColorTests(unittest.TestCase):
    def test_white(self):
        self.assertEqual(Color.White, tutor.search.color.parse("W"))

    def test_black(self):
        self.assertEqual(Color.Black, tutor.search.color.parse("B"))

    def test_blue(self):
        self.assertEqual(Color.Blue, tutor.search.color.parse("U"))

    def test_green(self):
        self.assertEqual(Color.Green, tutor.search.color.parse("G"))

    def test_red(self):
        self.assertEqual(Color.Red, tutor.search.color.parse("R"))

    def test_colorless(self):
        self.assertEqual(Color.Colorless, tutor.search.color.parse("C"))

    def test_ravnican_guilds(self):
        for guild, colors in guilds.items():
            with self.subTest(guild=guild):
                self.assertEqual(colors, tutor.search.guild.parse(guild))

    def test_alaran_shards(self):
        for shard, colors in shards.items():
            with self.subTest(shard=shard):
                self.assertEqual(colors, tutor.search.shard.parse(shard))


class SearchTests(unittest.TestCase):
    @given(s.text())
    def test_any_text(self, s):
        tutor.search.search.parse(s)

    @given(searches())
    def test_search(self, s):
        self.maxDiff = None
        query, expected = s
        self.assertEqual(str(expected), str(tutor.search.search.parse(query)))