def testTaxiPipelineCheckDagConstruction(self):
airflow_config = {
'schedule_interval': None,
'start_date': datetime.datetime(2019, 1, 1),
}
# Create directory structure and write expected user module file.
os.makedirs(os.path.join(self._test_dir, 'taxi'))
module_file = os.path.join(self._test_dir, 'taxi/taxi_utils.py')
with open(module_file, 'w') as f:
f.write('# Placeholder user module file.')
# Patch $HOME directory for pipeline DAG construction.
original_home = os.environ['HOME']
os.environ['HOME'] = self._test_dir
from tfx.examples.chicago_taxi_pipeline import taxi_pipeline_simple # pylint: disable=g-import-not-at-top
os.environ['HOME'] = original_home
logical_pipeline = taxi_pipeline_simple._create_pipeline(
pipeline_name='Test',
pipeline_root=self._test_dir,
data_root=self._test_dir,
module_file=module_file,
serving_model_dir=self._test_dir,
metadata_path=self._test_dir,
beam_pipeline_args=[])
self.assertEqual(9, len(logical_pipeline.components))
pipeline = AirflowDagRunner(
AirflowPipelineConfig(airflow_config)).run(logical_pipeline)
self.assertIsInstance(pipeline, models.DAG)
def test_taxi_pipeline_check_dag_construction(self):
airflow_config = {
'schedule_interval': None,
'start_date': datetime.datetime(2019, 1, 1),
}
logical_pipeline = taxi_pipeline_simple._create_pipeline()
self.assertEqual(9, len(logical_pipeline.components))
pipeline = TfxRunner(airflow_config).run(logical_pipeline)
self.assertIsInstance(pipeline, AirflowPipeline)
def testTaxiPipelineCheckDagConstruction(self):
airflow_config = {
'schedule_interval': None,
'start_date': datetime.datetime(2019, 1, 1),
}
logical_pipeline = taxi_pipeline_simple._create_pipeline(
pipeline_name='Test',
pipeline_root=self._test_dir,
data_root=self._test_dir,
module_file=self._test_dir,
serving_model_dir=self._test_dir,
metadata_path=self._test_dir)
self.assertEqual(9, len(logical_pipeline.components))
pipeline = AirflowDagRunner(airflow_config).run(logical_pipeline)
self.assertIsInstance(pipeline, models.DAG)