def test_unconnected_limited_infection(self):
"""
Tests the limited infection algorithm with a group of 3 unconnected users.
"""
u1 = User("u1", 1)
u2 = User("u2", 1)
u3 = User("u3", 1)
network = UserNetwork()
network.add_user(u1)
network.add_user(u2)
network.add_user(u3)
network.limited_infection(u1, 3, 2)
self.assertEqual(u1.version, 2)
self.assertEqual(u2.version, 2)
self.assertEqual(u3.version, 2)
network.limited_infection(u2, 3, 3)
self.assertEqual(u1.version, 3)
self.assertEqual(u2.version, 3)
self.assertEqual(u3.version, 3)
network.limited_infection(u3, 3, 4)
self.assertEqual(u1.version, 4)
self.assertEqual(u2.version, 4)
self.assertEqual(u3.version, 4)
network.limited_infection(u2, 2, 5)
self.assertEqual(u2.version, 5)
# One or other but not both
self.assertNotEqual(u1.version, u3.version)
self.assertTrue(u1.version == 5 or u3.version == 5)
def test_basic_limited_infection(self):
"""
Tests the limited infection algorithm on a network of 3 users such that user 1 coaches user 2 and
user 2 coaches user 3.
"""
u1 = User("u1", 1)
u2 = User("u2", 1)
u3 = User("u3", 1)
u1.add_student(u2)
u2.add_student(u3)
network = UserNetwork()
network.add_user(u1)
network.add_user(u2)
network.add_user(u3)
network.limited_infection(u1, 2, 2)
self.assertEqual(u1.version, 2)
self.assertEqual(u2.version, 2)
self.assertEqual(u3.version, 1)
network.limited_infection(u2, 2, 3)
self.assertEqual(u1.version, 2)
self.assertEqual(u2.version, 3)
self.assertEqual(u3.version, 3)
# Make sure it finds user w/ most infected students
network.limited_infection(u3, 2, 4)
self.assertEqual(u1.version, 2)
self.assertEqual(u2.version, 4)
self.assertEqual(u3.version, 4)
