test clean up

test/integration_tests/test_pytorch.py

@@ -16,28 +16,23 @@
 
 from src.catalog.catalog_manager import CatalogManager
 from test.util import perform_query
+from test.util import populate_catalog_with_built_in_udfs
 
 
 class PytorchTest(unittest.TestCase):
 
     def setUp(self):
         CatalogManager().reset()
+        populate_catalog_with_built_in_udfs()
 
     def test_should_run_pytorch_and_fastrcnn(self):
         query = """LOAD DATA INFILE 'data/ua_detrac/ua_detrac.mp4'
                    INTO MyVideo;"""
         perform_query(query)
 
-        create_udf_query = """CREATE UDF FastRCNNObjectDetector
-                  INPUT  (Frame_Array NDARRAY (3, 256, 256))
-                  OUTPUT (labels (10))
-                  TYPE  Classification
-                  IMPL  'src/udfs/classifier_udfs/fastrcnn_object_detector.py';
-        """
-        perform_query(create_udf_query)
-
         select_query = """SELECT FastRCNNObjectDetector(data) FROM MyVideo
                         WHERE id < 5;"""
+
         actual_batch = perform_query(select_query)
         self.assertEqual(actual_batch.batch_size, 5)
 

test/parser/test_parser_visitor.py

@@ -31,9 +31,8 @@ class ParserVisitorTests(unittest.TestCase):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-    def test_should_query_specification_visitor(self):
-        ParserVisitor.visit = MagicMock()
-        mock_visit = ParserVisitor.visit
+    @mock.patch.object(ParserVisitor, 'visit')
+    def test_should_query_specification_visitor(self, mock_visit):
         mock_visit.side_effect = ["columns",
                                   {"from": ["tables"], "where": "predicates"}]
 
@@ -160,17 +159,15 @@ def test_logical_expression(self):
         expected = visitor.visitLogicalExpression(ctx)
         self.assertEqual(expected, None)
 
-    def test_visit_string_literal_none(self):
+    @mock.patch.object(ParserVisitor, 'visitChildren')
+    def test_visit_string_literal_none(self, mock_visit):
         ''' Testing when string literal is None
             Function: visitStringLiteral
         '''
         visitor = ParserVisitor()
         ctx = MagicMock()
         ctx.STRING_LITERAL.return_value = None
 
-        ParserVisitor.visitChildren = MagicMock()
-        mock_visit = ParserVisitor.visitChildren
-
         visitor.visitStringLiteral(ctx)
         mock_visit.assert_has_calls([call(ctx)])
 
@@ -191,8 +188,6 @@ def test_visit_query_specification_base_exception(self):
         ''' Testing Base Exception error handling
             Function: visitQuerySpecification
         '''
-        ParserVisitor.visit = MagicMock()
-        ParserVisitor.visit
 
         visitor = ParserVisitor()
         ctx = MagicMock()

test/spark/test_session.py

@@ -27,10 +27,6 @@ def __init__(self, *args, **kwargs):
     def setUp(self):
         self.session = Session()
 
-    def tearDown(self):
-        self.session = Session()
-        self.session.stop()
-
     def test_session(self):
 
         spark_session = self.session.get_session()

test/udfs/ndarray_udfs/test_unnest.py

@@ -19,14 +19,14 @@
 from src.models.storage.batch import Batch
 from test.util import perform_query
 from test.util import populate_catalog_with_built_in_udfs
-from test.util import create_sample_video, NUM_FRAMES
 
 
-class UnnestTest(unittest.TestCase):
+class UnnestTests(unittest.TestCase):
+
     def setUp(self):
         CatalogManager().reset()
         populate_catalog_with_built_in_udfs()
-        create_sample_video(NUM_FRAMES)
+
         load_query = """LOAD DATA INFILE 'data/ua_detrac/ua_detrac.mp4'
                         INTO MyVideo;"""
         perform_query(load_query)
@@ -60,3 +60,7 @@ def test_should_unnest_dataframe(self):
         unnest_batch = perform_query(query)
         expected = Batch(self.unnest(without_unnest_batch.frames))
         self.assertEqual(unnest_batch, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/util.py

@@ -94,7 +94,9 @@ def create_dummy_batches(num_frames=NUM_FRAMES,
 
 def perform_query(query):
     stmt = Parser().parse(query)[0]
+    print(stmt)
     l_plan = StatementToPlanConvertor().visit(stmt)
+    print(l_plan)
     p_plan = PlanGenerator().build(l_plan)
     return PlanExecutor(p_plan).execute_plan()