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)
