class TestChoose(unittest.TestCase):
def setUp(self):
self.one = Choose({1})
self.six = Choose({1, 2, 3, 4, 5, 6})
self.repeated = Choose([1, 1, 1, 2])
def test_choose(self):
assert_almost_equal(self, self.one, Constant(1))
self.assertAlmostEqual(self.one.probability(1), 1.)
self.assertAlmostEqual(self.one.probability(0), 0.)
categorical_six = Categorical({x: 1 / 6 for x in range(1, 7)})
assert_almost_equal(self, self.six, categorical_six)
self.assertAlmostEqual(self.six.probability(1), 1 / 6)
self.assertAlmostEqual(self.six.probability(0), 0.)
def test_repeated(self):
counts = Counter(self.repeated.sample_n(1000))
self.assertLess(abs(counts[1] - 750), 50)
self.assertLess(abs(counts[2] - 250), 50)
table = self.repeated.table()
self.assertAlmostEqual(table[1], 0.75)
self.assertAlmostEqual(table[2], 0.25)
counts = Counter(self.repeated.sample_n(1000))
self.assertLess(abs(counts[1] - 750), 50)
self.assertLess(abs(counts[2] - 250), 50)
class TestChoose(unittest.TestCase):
def setUp(self):
self.one = Choose({1})
self.six = Choose({1, 2, 3, 4, 5, 6})
def test_choose(self):
assert_almost_equal(self, self.one, Constant(1))
self.assertAlmostEqual(self.one.probability(1), 1.)
self.assertAlmostEqual(self.one.probability(0), 0.)
categorical_six = Categorical({x: 1 / 6 for x in range(1, 7)})
assert_almost_equal(self, self.six, categorical_six)
self.assertAlmostEqual(self.six.probability(1), 1 / 6)
self.assertAlmostEqual(self.six.probability(0), 0.)
