import unittest from collections import namedtuple AssignedJob = namedtuple("AssignedJob", ["worker", "started_at"]) from sources.job_queue import assign_jobs, assign_jobs_improved class MyTestCase(unittest.TestCase): def test(self): num_workers = 2 jobs = [1, 2, 3, 4, 5] result = [AssignedJob(0, 0), AssignedJob(1, 0), AssignedJob(0, 1), AssignedJob(1, 2), AssignedJob(0, 4)] self.assertListEqual(result, assign_jobs_improved(num_workers, jobs)) def test1(self): num_workers = 4 jobs = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] result = [AssignedJob(0, 0), AssignedJob(1, 0), AssignedJob(2, 0), AssignedJob(3, 0), AssignedJob(0, 1), AssignedJob(1, 1), AssignedJob(2, 1), AssignedJob(3, 1), AssignedJob(0, 2), AssignedJob(1, 2), AssignedJob(2, 2), AssignedJob(3, 2), AssignedJob(0, 3), AssignedJob(1, 3), AssignedJob(2, 3), AssignedJob(3, 3), AssignedJob(0, 4), AssignedJob(1, 4), AssignedJob(2, 4), AssignedJob(3, 4)] self.assertListEqual(result, assign_jobs_improved(num_workers, jobs)) if __name__ == '__main__': unittest.main()