test_next_num.py

from next_num import RandomGen

import math
import unittest

# I would have used pytest for writing tests
# as it provides me decorators like parametrize
# but for the sake of the assignment I'm just sticking 
# with the standard library :)

class TestRandomGen(unittest.TestCase):

    def setUp(self):
        parameters = (
            # (random_nums, probabilities)
            ([1], [1]),
            ([4, 6, 1], [0, 0, 1]),
            ([3, 2, 6, 7], [0, 0.3, 0.5, 0.2]),
            ([-1, 0, 1, 2, 3], [0.01, 0.3, 0.58, 0.1, 0.01]),
        )

        self.random_gens = (RandomGen(rnums, probs) for rnums, probs in parameters)

    def test_result_in_random_nums(self):
        for rg in self.random_gens:
            rnum = rg.next_num()
            self.assertIn(rnum, rg._random_nums)

    def test_frequencies(self):
        for rg in self.random_gens:
            required_frequencies = dict(zip(rg._random_nums, rg._probabilities))

            observed_frequencies = {}
            total_runs = 100000

            for _ in range(total_runs):
                num = rg.next_num()
                observed_frequencies[num] = observed_frequencies.get(num, 0) + 1

            for num, req_freq in required_frequencies.items():
                observed_frequency = observed_frequencies.get(num, 0) / total_runs
                test = math.isclose(observed_frequency, req_freq, abs_tol=0.01)
                self.assertTrue(test, f"expected: {req_freq}, got: {observed_frequency}")


unittest.main()