def test_pool():
"""
Take a bucket of DRVs, and consider the results irrespective of order.
"""
pool = pools.pool(d6, count=2)
assert p(pool == pools.PlainResult(1, 1)) == Fraction(1, 36)
assert p(pool == pools.PlainResult(1, 2)) == Fraction(2, 36)
def test_pool_addition():
"""
You can add a constant, DRV or pool to a pool, and the effect is of
including one or more extra dice in the pool.
"""
pool = pools.pool(d6)
assert p(pool + 1 == pools.PlainResult(1, 1)) == Fraction(1, 6)
assert p(pool + d6 == pools.PlainResult(1, 1)) == Fraction(1, 36)
assert p(pool + pool == pools.PlainResult(1, 1)) == Fraction(1, 36)
def test_mixed_pool():
"""
Not all dice in pool need to be the same, and you can build up a pool one
item at a time if you want to.
"""
pool = pools.pool(d6, d8)
assert p(pool == pools.PlainResult(1, 1)) == Fraction(1, 48)
assert p(pool == pools.PlainResult(1, 2)) == Fraction(2, 48)
assert p(pool == pools.PlainResult(6, 7)) == Fraction(1, 48)
assert pool.is_same(pools.pool(d6) + d8)
def test_keep_lowest():
"""
Roll N, keep the worst K of some DRV.
"""
pool = pools.keep_lowest(2, d6, count=3)
assert p(pool == pools.PlainResult(6, 6)) == Fraction(1, 216)
# There are three ways each to get 1, 1, x for x = 2..6, plus 1, 1, 1.
assert p(pool == pools.PlainResult(1, 1)) == Fraction(16, 216)
pool0 = pools.keep_lowest(0, d6, count=10)
assert pool0.is_same(DRV({pools.PlainResult(): 1}))
def test_empty_pool():
"""
An empty pool has one possible value: the empty collection of values.
"""
empty1 = pools.pool()
empty2 = pools.pool(d6, count=0)
assert empty1.is_same(empty2)
assert empty1.to_dict() == {pools.PlainResult(): 1}
def test_result_repr(size):
"""
Because we have to do some work to set a sensible name for the subclasses
of PlainResult, we might as well test that.
"""
values = [1] * size
assert repr(pools.PlainResult(*values)).startswith('PlainResult(')
values = [0] * 20
assert repr(pools.KeepHighest(size)(*values)).startswith(f'Highest_{size}')
assert repr(pools.KeepLowest(size)(*values)).startswith(f'Lowest_{size}')
def test_keep_highest():
"""
Roll N, keep the best K of some DRV.
"""
pool = pools.keep_highest(2, d6, count=3)
# There are three ways each to get 6, 6, x for x = 1..5, plus 6, 6, 6.
assert p(pool == pools.PlainResult(6, 6)) == Fraction(16, 216)
assert p(pool == pools.PlainResult(1, 1)) == Fraction(1, 216)
# count=1000 acts as a performance test: if the implementation tries to
# compute all possibilities and then restrict to 0 dice, it will fail.
pool0 = pools.keep_highest(0, d6, count=1000)
assert pool0.is_same(DRV({pools.PlainResult(): 1}))
# Examples from docs
poolA = pools.keep_highest(2, d6) + d6 + d6
poolB = pools.pool(d6, count=3)
poolC = pools.keep_highest(2, d6, count=3)
poolD = pools.pool(d6, result_type=pools.KeepHighest(2)) + d6 + d6
assert poolA.is_same(poolB)
assert not poolA.is_same(poolC)
assert poolD.is_same(poolC)
def test_plain_result():
assert pools.PlainResult(1, 2) == pools.PlainResult(2, 1)
assert pools.PlainResult(1, 1, 2) == pools.PlainResult(1, 2, 1)
assert pools.PlainResult(1, 1) != pools.PlainResult(1, 2)
assert pools.PlainResult(1, 1, 2) != pools.PlainResult(1, 2, 2)
assert pools.PlainResult() != pools.PlainResult(0)
assert pools.PlainResult(1, 2) != (1, 2)
assert not pools.PlainResult(1, 2) == (1, 2)