def test_add_coaches(self):
one = User(0)
two = User(0)
one.add_coaches(two)
self.assertIn(two, one.coaches)
self.assertIn(one, two.coached_by)
def test_one_classroom_two_users_chain(self):
a = User()
b = User()
a.add_coaches(b)
infected_classrooms = total_infection(a)
self.assertEqual(infected_classrooms.infected, set([a, b]))
def test_spread_via_coaches_relationship(self):
a = User()
b = User()
a.add_coaches(b)
infected_classrooms = total_infection(a)
self.assertEqual(infected_classrooms.infected, set([a, b]))
def test_limit_1_two_users_chain(self):
a = User()
b = User()
a.add_coaches(b)
infection = LimitedInfection()
infected_classrooms = infection.limited_infection(a, 1)
self.assertEqual(infected_classrooms.infected, set([a]))
def test_one_classroom_three_users_hierarchical(self):
a = User()
b = User()
c = User()
a.add_coaches(b)
a.add_coaches(c)
infected_classrooms = total_infection(a)
self.assertEqual(infected_classrooms.infected, set([a, b, c]))
def test_limit_2_two_users_cyclical(self):
a = User()
b = User()
a.add_coaches(b)
b.add_coaches(a)
infection = LimitedInfection()
infected_classrooms = infection.limited_infection(a, 2)
self.assertEqual(infected_classrooms.infected, set([a, b]))
def test_limit_1_three_users_hierarchical(self):
a = User()
b = User()
c = User()
a.add_coaches(b)
a.add_coaches(c)
infection = LimitedInfection()
infected_classrooms = infection.limited_infection(a, 1)
self.assertEqual(infected_classrooms.infected, set([a]))
def test_limit_3_three_users_chain(self):
a = User()
b = User()
c = User()
a.add_coaches(b)
b.add_coaches(c)
infection = LimitedInfection()
infected_classrooms = infection.limited_infection(a, 3)
self.assertEqual(infected_classrooms.infected, set([a, b, c]))
def test_multiple_classrooms(self):
a = User()
b = User()
c = User()
d = User()
e = User()
a.add_coaches(b)
a.add_coaches(c)
d.add_coaches(e)
infected_classrooms = total_infection(d)
self.assertEqual(infected_classrooms.infected, set([d, e]))
def test_multiple_classrooms(self):
a = User()
b = User()
c = User()
d = User()
e = User()
a.add_coaches(b)
a.add_coaches(c)
d.add_coaches(e)
infection = LimitedInfection()
infected_classrooms = infection.limited_infection(b, 2)
self.assertEqual(infected_classrooms.infected, set([b, a]))
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]))])
