def test_execute(self, mock_hook): operator = KylinCubeOperator( task_id='kylin_task', dag=self.dag, **self._config ) hook = MagicMock() hook.invoke_command = ['fullbuild', 'build', 'merge', 'refresh', 'delete', 'build_streaming', 'merge_streaming', 'refresh_streaming', 'disable', 'enable', 'purge', 'clone', 'drop'] mock_hook.return_value = hook mock_hook.cube_run.return_value = {} self.assertIsNotNone(operator) self.assertEqual(self._config['kylin_conn_id'], operator.kylin_conn_id) self.assertEqual(self._config['project'], operator.project) self.assertEqual(self._config['cube'], operator.cube) self.assertEqual(self._config['command'], operator.command) self.assertEqual(self._config['start_time'], operator.start_time) self.assertEqual(self._config['end_time'], operator.end_time) operator.execute(None) mock_hook.assert_called_once_with( kylin_conn_id=self._config['kylin_conn_id'], project=self._config['project'], dsn=None ) mock_hook.return_value.cube_run.assert_called_once_with('kylin_sales_cube', 'build', end=datetime(2012, 1, 3, 0, 0), name=None, offset_end=None, offset_start=None, start=datetime(2012, 1, 2, 0, 0))
def test_execute_build(self, mock_hook): operator = KylinCubeOperator( is_track_job=True, timeout=5, interval=1, task_id='kylin_task', dag=self.dag, **self._config ) hook = MagicMock() hook.invoke_command = self.cube_command hook.cube_run.return_value = self.build_response hook.get_job_status.side_effect = ["RUNNING", "RUNNING", "FINISHED"] mock_hook.return_value = hook self.assertEqual(operator.execute(None)['uuid'], "c143e0e4-ac5f-434d-acf3-46b0d15e3dc6")
def test_execute_build_time_out_error(self, mock_hook): operator = KylinCubeOperator(is_track_job=True, timeout=5, interval=1, task_id='kylin_task', dag=self.dag, **self._config) hook = MagicMock() hook.invoke_command = self.cube_command hook.cube_run.return_value = self.build_response hook.get_job_status.return_value = "RUNNING" mock_hook.return_value = hook with pytest.raises(AirflowException): operator.execute(None)
def test_render_template(self): operator = KylinCubeOperator( task_id="kylin_build_1", kylin_conn_id='kylin_default', project="{{ params.project }}", cube="{{ params.cube }}", command="{{ params.command }}", start_time="{{ params.start_time }}", end_time="{{ params.end_time }}", is_track_job=True, dag=self.dag, params={ 'project': 'learn_kylin', 'cube': 'kylin_sales_cube', 'command': 'build', 'start_time': '1483200000000', 'end_time': '1483286400000', }, ) ti = TaskInstance(operator, DEFAULT_DATE) ti.render_templates() self.assertEqual('learn_kylin', getattr(operator, 'project')) self.assertEqual('kylin_sales_cube', getattr(operator, 'cube')) self.assertEqual('build', getattr(operator, 'command')) self.assertEqual('1483200000000', getattr(operator, 'start_time')) self.assertEqual('1483286400000', getattr(operator, 'end_time'))
def gen_build_time(): """ Gen build time and push to XCom (with key of "return_value") :return: A dict with build time values. """ return {'date_start': '1325347200000', 'date_end': '1325433600000'} gen_build_time_task = gen_build_time() gen_build_time_output_date_start = gen_build_time_task['date_start'] gen_build_time_output_date_end = gen_build_time_task['date_end'] build_task1 = KylinCubeOperator( task_id="kylin_build_1", command='build', start_time=gen_build_time_output_date_start, end_time=gen_build_time_output_date_end, is_track_job=True, dag=dag, ) build_task2 = KylinCubeOperator( task_id="kylin_build_2", command='build', start_time=gen_build_time_output_date_end, end_time='1325520000000', is_track_job=True, dag=dag, ) refresh_task1 = KylinCubeOperator( task_id="kylin_refresh_1",
""" ti = kwargs['ti'] ti.xcom_push(key='date_start', value='1325347200000') ti.xcom_push(key='date_end', value='1325433600000') gen_build_time_task = PythonOperator(python_callable=gen_build_time, task_id='gen_build_time', dag=dag) build_task1 = KylinCubeOperator( task_id="kylin_build_1", kylin_conn_id='kylin_default', project='learn_kylin', cube='kylin_sales_cube', command='build', start_time=gen_build_time_task.output['date_start'], end_time=gen_build_time_task.output['date_end'], is_track_job=True, dag=dag, ) build_task2 = KylinCubeOperator( task_id="kylin_build_2", kylin_conn_id='kylin_default', project='learn_kylin', cube='kylin_sales_cube', command='build', start_time='1325433600000', end_time='1325520000000', is_track_job=True,