class TestGraph(TestCase): def setUp(self): init_db() self.graph = Graph(db_session) def tearDown(self): Base.metadata.drop_all(bind=engine) @classmethod def tearDownClass(cls): db_session.remove() def test_add(self): self.graph.add(actor_item) self.assertEqual(self.graph.get_actor("a").wiki_page, actor_item.get("wiki_page"))
class TestAnalysis(TestCase): def setUp(self): init_db() self.graph = Graph(db_session) def tearDown(self): Base.metadata.drop_all(bind=engine) @classmethod def tearDownClass(cls): db_session.remove() def test_hub_actor(self): self.graph.add(actor_item) self.graph.add(movie_item) self.graph.add_actor({"name": "foo"}) actors = self.graph.get_hub_actor() self.assertEqual(actors[0][0].name, "a") actors2 = self.graph.get_hub_actor(n=-1) self.assertEqual(len(actors), len(actors2)) self.graph.get_hub_actor(plot=True, save_to="output/hub_actor.png") def test_age_correlation(self): self.graph.add(actor_item) self.graph.add(movie_item) self.graph.add_actor({"name": "foo", "wiki_page": "b", "age": 10}) self.graph.add_actor({"name": "bar", "age": 20, "total_gross": 12345}) actors = self.graph.get_age_correlation() self.assertEqual(actors[0][0], 20) self.graph.get_age_correlation(plot=True, save_to="output/age_correlation.png")