def test_connect_users_in_a_graph(self):
graph = Graph()
user1 = User()
user2 = User()
graph.add_user(user1)
graph.add_user(user2)
graph.connect_users(user1.get_id(), user2.get_id())
assert user2 in user1.get_connections() and user1 in user2.get_connections()
def construct_graph_of_count(count):
graph = Graph()
user1 = User()
graph.add_user(user1)
for i in range(0, count-1):
user = User()
graph.add_user(user)
graph.connect_users(user1.get_id(), user.get_id())
return graph
class TestTotalInfection(unittest.TestCase):
def setUp(self):
self.graph = Graph()
self.user1 = User()
self.user2 = User()
self.user3 = User()
self.user4 = User()
self.users = [self.user1, self.user2, self.user3, self.user4]
def test_total_infection_on_component_with_one_user(self):
user = self.user1
self.graph.add_user(user)
infection_count = infect_total(user)
assert infection_count == 1
verify_users_in_a_group_are_infected(self.graph)
def test_total_infection_on_component_with_two_users(self):
user1 = self.user1
user2 = self.user2
self.graph.add_user(user1)
self.graph.add_user(user2)
self.graph.connect_users(user1.get_id(), user2.get_id())
infection_count = infect_total(user1)
assert infection_count == 2
verify_users_in_a_group_are_infected(self.graph)
def test_total_infection_on_component_with_three_users_in_a_cycle(self):
user1 = self.user1
user2 = self.user2
user3 = self.user3
users = [user1, user2, user3]
for user in users:
self.graph.add_user(user)
self.graph.connect_users(user1.get_id(), user2.get_id())
self.graph.connect_users(user1.get_id(), user3.get_id())
self.graph.connect_users(user2.get_id(), user3.get_id())
infection_count = infect_total(user1)
assert infection_count == 3
verify_users_in_a_group_are_infected(self.graph)
def test_total_infection_on_components_with_three_users_connected_to_one_user(self):
user1 = self.user1
self.graph.add_user(user1)
users = self.users
for user in users[1:]:
self.graph.add_user(user)
self.graph.connect_users(user1.get_id(), user.get_id())
infection_count = infect_total(users[1])
assert infection_count == 4
verify_users_in_a_group_are_infected(self.graph)