def test_add_uninfected_user(self):
classrooms = Classrooms()
user = User(0)
classrooms.add_user(user)
self.assertIn(user, classrooms.uninfected)
def limited_infection(self, starting_user, limit):
classrooms = Classrooms()
def limited_infect(user):
# Make sure limit isn't reached.
if user.version == 0 and self.count < limit:
# Moving a user from the uninfected set to the infected set.
user.infect()
self.count += 1 # Increment counter.
classrooms.add_user(user)
# All relationships (union of coached_by and coaches sets).
relationships = user.coached_by | user.coaches
for relationship in relationships:
limited_infect(relationship)
limited_infect(starting_user)
# Used to ensure we have users left to infect.
remaining = classrooms.total_users() - limit
if remaining > 0:
while self.count <= limit:
# Take a user that is uninfected and infect it's "classroom".
limited_infect(classrooms.uninfected.pop())
return classrooms
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_total_users(self):
classrooms = Classrooms()
uninfected_user = User(0)
infected_user = User(1)
classrooms.add_user(uninfected_user)
classrooms.add_user(infected_user)
self.assertEqual(classrooms.total_users(), 2)
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]))])
