class TestTree(unittest.TestCase): def setUp(self): self.tournament = Tournament() self.teams = self.tournament.teams self.tree = ProbableTournamentTree(self.tournament) def test_get_opponent_at_stage(self): brazil = Team.get_for_country(self.teams, 'Brazil') mexico = Team.get_for_country(self.teams, 'Mexico') cameroon = Team.get_for_country(self.teams, 'Cameroon') spain = Team.get_for_country(self.teams, 'Spain') netherlands = Team.get_for_country(self.teams, 'Netherlands') opp = self.tree.get_opponent_at_stage(brazil, 0) self.assertEqual(opp.country, netherlands.country) opp = self.tree.get_opponent_at_stage(brazil, 1) self.assertEqual(opp.country, Team.get_for_country(self.teams, 'Colombia').country) opp = self.tree.get_opponent_at_stage(brazil, 3) self.assertEqual(opp.country, spain.country) opp = self.tree.get_opponent_at_stage(netherlands, 0) self.assertEqual(opp.country, brazil.country) opp = self.tree.get_opponent_at_stage(mexico, 0) self.assertEqual(opp.country, spain.country) # test for a team that isn't in the probability tree self.assertEqual(self.tree.get_opponent_at_stage(cameroon, 0).country, self.tree.get_opponent_at_stage(mexico, 0).country)
def setUp(self): self.tournament = Tournament() self.teams = self.tournament.teams self.tree = ProbableTournamentTree(self.tournament)