def test_add_users_to_a_graph(self): graph = Graph() user = User() user_id = user.get_id() graph.add_user(user) assert user_id in graph.users and graph.users[user_id] == user
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)