Fixed tests and added CI (#5)

TheMichaelHu · web-flow · commit 17f3512e1837 · 2020-06-05T11:30:03.000-04:00
* add dockerfile for testing
* add cloud build and travis CI
* fixed tests
* add python version badge
diff --git a/.travis.yml b/.travis.yml
@@ -0,0 +1,9 @@
+language: python
+python:
+  - "3.6"
+  - "3.7"
+  # Tensorflow 1.x does not support python 3.8+
+install:
+  - pip install -e .
+script:
+  - python -m unittest
diff --git a/Dockerfile b/Dockerfile
@@ -0,0 +1,19 @@
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+FROM python:3.6
+COPY . /app
+WORKDIR /app
+RUN pip install -e ".[dev]"
+RUN python -m unittest
diff --git a/README.md b/README.md
@@ -1,5 +1,8 @@
 # ML Pipeline Generator
+![PyPI - Python Version](https://img.shields.io/pypi/pyversions/ml-pipeline-gen)
 [![PyPI version](https://badge.fury.io/py/ml-pipeline-gen.svg)](https://badge.fury.io/py/ml-pipeline-gen)
+[![Build
+Status](https://travis-ci.com/GoogleCloudPlatform/ml-pipeline-generator-python.svg?branch=master)](https://travis-ci.com/GoogleCloudPlatform/ml-pipeline-generator-python)
 
 ML Pipeline Generator is a tool for generating end-to-end pipelines composed of GCP components so that users can easily migrate their local ML models onto GCP and start realizing the benefits of the Cloud quickly. 
 
@@ -131,5 +134,5 @@ as inputs in the config file.
 To modify the behavior of the library, install `ml-pipeline-gen` using:
 
 ```bash
-pip install -e .
+pip install -e ".[dev]"
 ```
diff --git a/setup.py b/setup.py
@@ -1,13 +1,13 @@
 # Copyright 2020 Google Inc. All Rights Reserved.
 #
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Licensed under the Apache License, Version 2.0 (the 'License');
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
+# distributed under the License is distributed on an 'AS IS' BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
@@ -16,37 +16,43 @@
 import setuptools
 import ml_pipeline_gen
 
-with open("README.md", "r") as f:
+with open('README.md', 'r') as f:
     long_description = f.read()
 
 setuptools.setup(
-    name="ml-pipeline-gen",
+    name='ml-pipeline-gen',
     version=ml_pipeline_gen.__version__,
-    author="Michael Hu",
-    author_email="author@example.com",
-    description="A tool for generating end-to-end pipelines on GCP.",
+    author='Michael Hu',
+    author_email='author@example.com',
+    description='A tool for generating end-to-end pipelines on GCP.',
     long_description=long_description,
-    long_description_content_type="text/markdown",
-    url="https://github.com/GoogleCloudPlatform/ml-pipeline-generator-python",
-    packages=["ml_pipeline_gen"],
+    long_description_content_type='text/markdown',
+    url='https://github.com/GoogleCloudPlatform/ml-pipeline-generator-python',
+    packages=['ml_pipeline_gen'],
     install_requires=[
-        "cloudml-hypertune",
-        "gcsfs",
-        "google-api-python-client",
-        "jinja2",
-        "joblib",
-        "kfp",
-        "pandas",
-        "pyyaml",
-        "scikit-learn",
-        "tensorflow>=0.14.0,<2.0.0",
-        "xgboost",
+        'cloudml-hypertune',
+        'gcsfs',
+        'google-api-python-client',
+        'jinja2',
+        'joblib',
+        'kfp',
+        'pandas',
+        'pyyaml',
+        'scikit-learn',
+        'tensorflow>=1.14.0,<2.0.0',
+        'xgboost',
     ],
+    extras_require={
+        'dev': [
+            'mock',
+        ]
+    },
     classifiers=[
-        "Programming Language :: Python :: 3",
-        "License :: OSI Approved :: Apache Software License",
-        "Operating System :: OS Independent",
+        'Programming Language :: Python :: 3.6',
+        'Programming Language :: Python :: 3.7',
+        'License :: OSI Approved :: Apache Software License',
+        'Operating System :: OS Independent',
     ],
-    python_requires=">=3.6",
+    python_requires='>=3.6',
     include_package_data=True,
 )
diff --git a/tests/examples/sklearn/test_sklearn_model.py b/tests/examples/sklearn/test_sklearn_model.py
@@ -12,29 +12,66 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Unit tests demo scikit-learn model."""
+import argparse
+import os
+import shutil
+import sys
+import tempfile
 import unittest
 
-from examples.sklearn import sklearn_model
+from tests import test_utils
 
 
 class TestModel(unittest.TestCase):
     """Tests demo model."""
 
+    @classmethod
+    def setUpClass(cls):
+        super(TestModel, cls).setUpClass()
+        cls.test_dir = tempfile.mkdtemp()
+        cls.demo_dir = os.path.join(cls.test_dir, 'demo')
+        shutil.copytree('examples/sklearn', cls.demo_dir)
+
+        # TODO(humichael) We can't import the model using __import__ because
+        # several other examples are also adding their demo dirs to sys.path.
+        # It's very likely the model module that is imported is not the one from
+        # this test. All examples currently use the same census_preprocess.
+        # These tests will break if any example uses a different preprocessing
+        # script.
+        # We should just mock this.
+        sys.path.append(cls.demo_dir)
+        sklearn_model = test_utils.load_module(
+            'sklearn_model', os.path.join(
+                cls.demo_dir, 'model', 'sklearn_model.py'))
+        sklearn_preprocess = test_utils.load_module(
+            'sklearn_preprocess', os.path.join(
+                cls.demo_dir, 'model', 'census_preprocess.py'))
+        sys.path.remove(cls.demo_dir)
+        params = argparse.Namespace(C=1.0)
+        cls.model = sklearn_model.get_model(params)
+        cls.features, cls.labels, _, _ = sklearn_preprocess.load_data()
+
+    @classmethod
+    def tearDownClass(cls):
+        super(TestModel, cls).tearDownClass()
+        shutil.rmtree(cls.test_dir)
+
     def setUp(self):
         super(TestModel, self).setUp()
-        self.features, self.labels = sklearn_model.get_data()
+        self.model = self.__class__.model
+        self.features = self.__class__.features
+        self.labels = self.__class__.labels
 
     def test_get_data(self):
         """Checks that there is a label for each feature."""
         self.assertEqual(self.features.shape[0], self.labels.shape[0])
 
     def test_get_model(self):
         """Checks that the model can be trained and used for predictions."""
-        model = sklearn_model.get_model()
-        model.fit(self.features, self.labels)
-        preds = model.predict(self.features)
-        self.assertEqual(preds.shape, self.labels.shape)
+        self.model.fit(self.features, self.labels)
+        preds = self.model.predict(self.features)
+        self.assertEqual(preds.shape[0], self.labels.shape[0])
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     unittest.main()
diff --git a/tests/examples/tensorflow/test_tf_model.py b/tests/examples/tensorflow/test_tf_model.py
@@ -12,29 +12,75 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Unit tests demo TF model."""
+import argparse
+import os
+import shutil
+import sys
+import tempfile
 import unittest
 
-from examples.tf import tf_model
+import tensorflow.compat.v1 as tf
 
+from tests import test_utils
 
-class TestModel(unittest.TestCase):
+
+class TestModel(tf.test.TestCase):
     """Tests TF demo model."""
 
+    @classmethod
+    def setUpClass(cls):
+        super(TestModel, cls).setUpClass()
+        cls.test_dir = tempfile.mkdtemp()
+        cls.demo_dir = os.path.join(cls.test_dir, 'demo')
+        shutil.copytree('examples/tf', cls.demo_dir)
+
+        # TODO(humichael) We can't import the model using __import__ because
+        # several other examples are also adding their demo dirs to sys.path.
+        # It's very likely the model module that is imported is not the one from
+        # this test. All examples currently use the same census_preprocess.
+        # These tests will break if any example uses a different preprocessing
+        # script.
+        # We should just mock this.
+        sys.path.append(cls.demo_dir)
+        tf_model = test_utils.load_module(
+            'tf_model', os.path.join(cls.demo_dir, 'model', 'tf_model.py'))
+        tf_preprocess = test_utils.load_module(
+            'tf_preprocess', os.path.join(
+                cls.demo_dir, 'model', 'census_preprocess.py'))
+        sys.path.remove(cls.demo_dir)
+
+        cls.features, cls.labels, _, _ = tf_preprocess.load_data()
+        cls.model = tf_model
+
+    @classmethod
+    def tearDownClass(cls):
+        super(TestModel, cls).tearDownClass()
+        shutil.rmtree(cls.test_dir)
+
+    # pylint: disable=g-import-not-at-top
     def setUp(self):
         super(TestModel, self).setUp()
-        self.features, self.labels = tf_model.get_data()
+        self.model = self.__class__.model
+        self.features = self.__class__.features
+        self.labels = self.__class__.labels
 
     def test_get_data(self):
         """Checks that there is a label for each feature."""
         self.assertEqual(self.features.shape[0], self.labels.shape[0])
 
     def test_get_model(self):
         """Checks that the model can be trained and used for predictions."""
-        model = tf_model.get_model()
+        input_layer = tf.keras.layers.Input(shape=(self.features.shape[1],))
+        params = argparse.Namespace(first_layer_size=50, num_layers=5)
+        predictions = self.model.get_model(input_layer, params)
+
+        model = tf.keras.models.Model(inputs=input_layer, outputs=predictions)
+        model.compile(optimizer='adam', loss=tf.losses.sigmoid_cross_entropy,
+                      metrics=['accuracy'])
         model.fit(self.features, self.labels)
         preds = model.predict(self.features)
         self.assertEqual(preds.shape[0], self.labels.shape[0])
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     unittest.main()
diff --git a/tests/ml_pipeline_gen/test_models.py b/tests/ml_pipeline_gen/test_models.py
@@ -13,9 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Unit tests for models classes."""
+import mock
 import os
+import shutil
+import tempfile
 import unittest
 
+from googleapiclient import discovery
+
 from ml_pipeline_gen.models import BaseModel
 from ml_pipeline_gen.models import SklearnModel
 
@@ -33,45 +38,54 @@ class TestSklearnModel(unittest.TestCase):
     """Tests SklearnModel class."""
 
     @classmethod
-    def setUpClass(cls):
-        """Instantiates a model."""
+    @mock.patch.object(discovery, 'build')
+    def setUpClass(cls, build_mock):
+        """Copies a demo and instantiates a model."""
         super(TestSklearnModel, cls).setUpClass()
-        cls.config = "examples/sklearn/config.yaml"
+        build_mock.return_value = None
+        cls.cwd = os.getcwd()
+        cls.test_dir = tempfile.mkdtemp()
+        cls.demo_dir = os.path.join(cls.test_dir, 'demo')
+        shutil.copytree('examples/sklearn', cls.demo_dir)
+
+        os.chdir(cls.demo_dir)
+        cls.config = 'config.yaml.example'
         cls.model = SklearnModel(cls.config)
 
     @classmethod
     def tearDownClass(cls):
-        """Cleans up generated directories."""
+        """Switch back to the original working dir and removes the demo."""
         super(TestSklearnModel, cls).tearDownClass()
-        cls.model.clean_up()
+        os.chdir(cls.cwd)
+        shutil.rmtree(cls.test_dir)
 
-    # TODO(humichael): technically private functions don't need to be tested. It
-    # should reflect in public functions.
-    def test_set_config(self):
-        """Ensures instance variables are created."""
-        model = self.__class__.model
-        model.model = {}
+    def setUp(self):
+        super(TestSklearnModel, self).setUp()
+        self.model = self.__class__.model
 
-        model._set_config(self.__class__.config)
-        self.assertEqual(model.model["name"], "sklearn_demo_model")
+    def tearDown(self):
+        super(TestSklearnModel, self).tearDown()
+        try:
+            self.__class__.model.clean_up()
+        except FileNotFoundError:
+            pass
 
-    def test_populate_trainer(self):
+    def test_generate_files(self):
         """Ensures task.py and model.py are created."""
-        model = self.__class__.model
-        model.clean_up()
-
-        model._populate_trainer()
-        trainer_files = os.listdir("trainer")
-        self.assertIn("task.py", trainer_files)
-        self.assertIn("model.py", trainer_files)
-
-    @unittest.skip("How to test without running training?")
+        self.assertFalse(os.path.exists('trainer'))
+        self.model.generate_files()
+        self.assertTrue(os.path.exists('trainer'))
+        trainer_files = os.listdir('trainer')
+        self.assertIn('task.py', trainer_files)
+        self.assertIn('model.py', trainer_files)
+
+    @unittest.skip('How to test without running training?')
     def test_local_train(self):
         """Tests local training."""
-        model = self.__class__.model
-        model.train()
-        model_files = os.listdir("models")
-        self.assertIn("{}.joblib".format(model.model["name"]), model_files)
+        self.model.generate_files()
+        self.model.train()
+        model_files = os.listdir('models')
+        self.assertIn('{}.joblib'.format(self.model.model['name']), model_files)
 
     # TODO(humichael): Need to spoof CAIP calls to test this.
     def test_cloud_train(self):
@@ -84,5 +98,5 @@ def test_serve(self):
         pass
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     unittest.main()
diff --git a/tests/test_utils.py b/tests/test_utils.py
