Skip to content

Commit

Permalink
Fix handler_factory_test to avoid kokoro failures.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 260254635
  • Loading branch information
tfx-copybara authored and tensorflow-extended-team committed Jul 27, 2019
1 parent ece969f commit c6deec0
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion tfx/tools/cli/handler/handler_factory_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from __future__ import print_function

import os
import sys
import tempfile

import mock
Expand All @@ -27,7 +28,6 @@
from tfx.tools.cli.handler import airflow_handler
from tfx.tools.cli.handler import beam_handler
from tfx.tools.cli.handler import handler_factory
from tfx.tools.cli.handler import kubeflow_handler


class _MockClientClass(object):
Expand All @@ -42,6 +42,7 @@ class HandlerFactoryTest(tf.test.TestCase):
def setUp(self):
super(HandlerFactoryTest, self).setUp()
self.flags_dict = {}
sys.modules['kfp'] = mock.Mock()

def test_create_handler_airflow(self):
self.flags_dict[labels.ENGINE_FLAG] = 'airflow'
Expand All @@ -57,6 +58,7 @@ def test_create_handler_kubeflow(self):
labels.IAP_CLIENT_ID: 'dummyID',
labels.NAMESPACE: 'kubeflow',
}
from tfx.tools.cli.handler import kubeflow_handler # pylint: disable=g-import-not-at-top
self.assertIsInstance(
handler_factory.create_handler(flags_dict),
kubeflow_handler.KubeflowHandler)
Expand Down

0 comments on commit c6deec0

Please sign in to comment.