def test_count_0_one_classroom(self): a = User() b = User() c = User() a.add_coaches(b) a.add_coaches(c) classrooms = Classrooms() classrooms.add_user(a) classrooms.add_user(b) classrooms.add_user(c) infection = LargestInfection() largest_infections = infection.find_largest_infection(classrooms, 0) self.assertEqual(largest_infections, [])
def test_count_2_two_classrooms(self): a = User() b = User() c = User() d = User() e = User() a.add_coaches(b) a.add_coaches(c) d.add_coaches(e) classrooms = Classrooms() classrooms.add_user(a) classrooms.add_user(b) classrooms.add_user(c) classrooms.add_user(d) classrooms.add_user(e) infection = LargestInfection() largest_infections = infection.find_largest_infection(classrooms, 2) self.assertEqual(largest_infections, [(3, set([a, b, c])), (2, set([d, e]))])