Refactor tests

twosigma · Jul 2, 2018 · b560b00 · b560b00
1 parent ce05ef1
commit b560b00
Show file tree

Hide file tree

Showing 14 changed files with 2,126 additions and 3,908 deletions.
diff --git a/python/tests/README.md b/python/tests/README.md
@@ -40,14 +40,6 @@ To run the tests issue the following command from the root Flint directory:
 scripts/run_python_tests.sh
 ```
 
-## Code
-The code for the tests is found in this `tests` directory. The content of the files here is as follows:
-
-* `base_test_case.py` Contains code for the `BaseTestCase` abstract class that is the grandfather of all the testcases.
-* `spark_test_case.py` Contains a concrete class, `SparkTestCase`, that inherits from `BaseTestCase` and sets up a local `SparkContext`. This is the default class to inherit test cases from.
-* `test_dataframe.py` Contains about 50 test cases for the `TimeSeriesDataFrame`.
-* `test_data.py` Contains constant data for the tests.
-* `utils.py` Contains specialized assert functions and Pandas DataFrame creation.
 
 ## Extending
 If the test setup done in the default class, `SparkTestCase`, does not fit the needs of a particular environment, a new class can be written. The name of the new class, say `MyTestCase` is then exported in the `BASE_CLASS` variable before the tests are run:

diff --git a/python/tests/ts/base_test_case.py → python/tests/base_test_case.py b/python/tests/ts/base_test_case.py → python/tests/base_test_case.py
@@ -19,7 +19,7 @@
 import unittest
 from abc import ABCMeta, abstractclassmethod
 import tests.utils as test_utils
-from tests.ts.test_data import (FORECAST_DATA, PRICE_DATA, VOL_DATA, VOL2_DATA,
+from tests.test_data import (FORECAST_DATA, PRICE_DATA, PRICE2_DATA, VOL_DATA, VOL2_DATA,
                              VOL3_DATA, INTERVALS_DATA)
 from functools import lru_cache
 
@@ -62,11 +62,12 @@ def price(self):
         return self.flintContext.read.pandas(
             test_utils.make_pdf(PRICE_DATA, ["time", "id", "price"]))
 
+    @lru_cache(maxsize=None)
+    def price2(self):
+        return self.flintContext.read.pandas(
+            test_utils.make_pdf(PRICE2_DATA, ["time", "id", "price"]))
+
     @lru_cache(maxsize=None)
     def intervals(self):
         return self.flintContext.read.pandas(
             test_utils.make_pdf(INTERVALS_DATA, ['time']))
-
-    def clocks(self):
-        from ts.flint import clocks
-        return clocks
diff --git a/python/tests/ts/spark_test_case.py → python/tests/spark_test_case.py b/python/tests/ts/spark_test_case.py → python/tests/spark_test_case.py
@@ -13,6 +13,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+
 ''' The common code for all Flint unit tests '''
 import os
 import sys
@@ -58,6 +59,10 @@ def __teardown(cls):
         cls.sc.stop()
         cls.sc._gateway.shutdown()
         cls.sc._gateway = None
+
+        SparkContext._jvm = None
+        SparkContext._gateway = None
+
         delattr(cls, 'sqlContext')
         delattr(cls, 'sc')
         os.environ = cls._env

diff --git a/python/tests/ts/test_data.py → python/tests/test_data.py b/python/tests/ts/test_data.py → python/tests/test_data.py
@@ -55,6 +55,13 @@
     (1250, 7, 6.0,)
 ]
 
+PRICE2_DATA = [
+    (0, 1, 1.0,),
+    (0, 2, 2.0,),
+    (1, 1, 3.0,),
+    (1, 2, 4.0,),
+    (1, 3, 5.0,),
+]
 
 VOL_DATA = [
     (1000, 7, 100,),

diff --git a/python/tests/ts/__init__.py b/python/tests/ts/__init__.py