def testExportToPipeline(self, mock_get_template):
self._setupTestNotebook()
c = interactive_context.InteractiveContext()
export_filepath = os.path.join(self._exportdir, 'exported_pipeline.py')
c.export_to_pipeline(notebook_filepath=self._notebook_fp,
export_filepath=export_filepath,
runner_type='beam')
with open(export_filepath, 'r') as exported_pipeline:
code = exported_pipeline.read()
self.assertEqual(
code,
textwrap.dedent('''\
foo = 1
def bar():
a = "hello"
b = "world"
return a + b
def baz():
c = "nyan"
d = "cat"
return c + d'''))
def testUnresolvedChannel(self):
class _FakeComponentSpec(types.ComponentSpec):
PARAMETERS = {}
INPUTS = {
'input':
component_spec.ChannelParameter(
type=standard_artifacts.Examples)
}
OUTPUTS = {}
class _FakeExecutor(base_executor.BaseExecutor):
CALLED = False
def Do(self, input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any]) -> None:
_FakeExecutor.CALLED = True
class _FakeComponent(base_component.BaseComponent):
SPEC_CLASS = _FakeComponentSpec
EXECUTOR_SPEC = executor_spec.ExecutorClassSpec(_FakeExecutor)
def __init__(self, spec: types.ComponentSpec):
super(_FakeComponent, self).__init__(spec=spec)
c = interactive_context.InteractiveContext()
foo = types.Channel(type=standard_artifacts.Examples,
artifacts=[standard_artifacts.Examples()])
component = _FakeComponent(_FakeComponentSpec(input=foo))
with self.assertRaisesRegexp(ValueError, 'Unresolved input channel'):
c.run(component)
def testBasicRun(self):
class _FakeComponentSpec(types.ComponentSpec):
PARAMETERS = {}
INPUTS = {}
OUTPUTS = {}
class _FakeExecutor(base_executor.BaseExecutor):
CALLED = False
def Do(self, input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any]) -> None:
_FakeExecutor.CALLED = True
class _FakeComponent(base_component.BaseComponent):
SPEC_CLASS = _FakeComponentSpec
EXECUTOR_SPEC = executor_spec.ExecutorClassSpec(_FakeExecutor)
def __init__(self, spec: types.ComponentSpec):
super(_FakeComponent, self).__init__(spec=spec)
c = interactive_context.InteractiveContext()
component = _FakeComponent(_FakeComponentSpec())
c.run(component)
self.assertTrue(_FakeExecutor.CALLED)
def testExportToPipelineRaisesErrorInvalidRunnerType(self):
self._setupTestNotebook()
c = interactive_context.InteractiveContext()
export_filepath = os.path.join(self._exportdir, 'exported_pipeline.py')
with self.assertRaisesRegexp(ValueError, 'runner_type'):
c.export_to_pipeline(notebook_filepath=self._notebook_fp,
export_filepath=export_filepath,
runner_type='foobar')
def testShow(self, *unused_mocks):
context = interactive_context.InteractiveContext()
mock_object = mock.MagicMock()
standard_visualizations.ExampleAnomaliesVisualization.display = mock_object
mock_object.assert_not_called()
artifact = standard_artifacts.ExampleAnomalies()
context.show(
types.Channel(type=standard_artifacts.ExampleAnomalies,
artifacts=[artifact]))
mock_object.assert_called_with(artifact)
def testTelemetry(self, mock_launcher_create):
class _FakeLauncher(object):
def __init__(self):
self.recorded_labels = []
def launch(self):
self.recorded_labels = telemetry_utils.make_beam_labels_args()
return mock.MagicMock()
class _FakeComponentSpec(types.ComponentSpec):
PARAMETERS = {}
INPUTS = {}
OUTPUTS = {}
class _FakeExecutor(base_executor.BaseExecutor):
def Do(self, input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any]) -> None:
pass
class _FakeComponent(base_component.BaseComponent):
SPEC_CLASS = _FakeComponentSpec
EXECUTOR_SPEC = executor_spec.ExecutorClassSpec(_FakeExecutor)
def __init__(self):
super(_FakeComponent, self).__init__(spec=_FakeComponentSpec())
# Set up fake on launcher.
fake_launcher = _FakeLauncher()
mock_launcher_create.side_effect = [fake_launcher]
InProcessComponentLauncher.create = mock_launcher_create
context = interactive_context.InteractiveContext()
context.run(_FakeComponent())
self.assertIn('--labels tfx_runner=interactivecontext',
' '.join(fake_launcher.recorded_labels))
def testRunMethodRequiresIPython(self):
del builtins.__dict__['__IPYTHON__']
c = interactive_context.InteractiveContext()
self.assertIsNone(c.run(None))
