class TestKylinHook(unittest.TestCase):
def setUp(self) -> None:
self.hook = KylinHook(kylin_conn_id='kylin_default',
project='learn_kylin')
@patch("kylinpy.Kylin.get_job")
def test_get_job_status(self, mock_job):
job = MagicMock()
job.status = "ERROR"
mock_job.return_value = job
self.assertEqual(self.hook.get_job_status('123'), "ERROR")
@patch("kylinpy.Kylin.get_datasource")
def test_cube_run(self, cube_source):
class MockCubeSource:
def invoke_command(self, command, **kwargs):
invoke_command_list = [
'fullbuild',
'build',
'merge',
'refresh',
'delete',
'build_streaming',
'merge_streaming',
'refresh_streaming',
'disable',
'enable',
'purge',
'clone',
'drop',
]
if command in invoke_command_list:
return {"code": "000", "data": {}}
else:
raise KylinCubeError(
f'Unsupported invoke command for datasource: {command}'
)
cube_source.return_value = MockCubeSource()
response_data = {"code": "000", "data": {}}
self.assertDictEqual(self.hook.cube_run('kylin_sales_cube', 'build'),
response_data)
self.assertDictEqual(self.hook.cube_run('kylin_sales_cube', 'refresh'),
response_data)
self.assertDictEqual(self.hook.cube_run('kylin_sales_cube', 'merge'),
response_data)
self.assertDictEqual(
self.hook.cube_run('kylin_sales_cube', 'build_streaming'),
response_data)
self.assertRaises(
AirflowException,
self.hook.cube_run,
'kylin_sales_cube',
'build123',
)
def execute(self, context):
_hook = KylinHook(kylin_conn_id=self.kylin_conn_id,
project=self.project,
dsn=self.dsn)
_support_invoke_command = kylinpy.CubeSource.support_invoke_command
if self.command.lower() not in _support_invoke_command:
raise AirflowException(
'Kylin:Command {} can not match kylin command list {}'.format(
self.command, _support_invoke_command))
kylinpy_params = {
'start':
datetime.fromtimestamp(int(self.start_time) /
1000) if self.start_time else None,
'end':
datetime.fromtimestamp(int(self.end_time) /
1000) if self.end_time else None,
'name':
self.segment_name,
'offset_start':
int(self.offset_start) if self.offset_start else None,
'offset_end':
int(self.offset_end) if self.offset_end else None,
}
rsp_data = _hook.cube_run(self.cube, self.command.lower(),
**kylinpy_params)
if self.is_track_job and self.command.lower() in self.build_command:
started_at = timezone.utcnow()
job_id = rsp_data.get("uuid")
if job_id is None:
raise AirflowException("kylin job id is None")
self.log.info("kylin job id: %s", job_id)
job_status = None
while job_status not in self.jobs_end_status:
if (timezone.utcnow() -
started_at).total_seconds() > self.timeout:
raise AirflowException(
'kylin job {} timeout'.format(job_id))
time.sleep(self.interval)
job_status = _hook.get_job_status(job_id)
self.log.info('Kylin job status is %s ', job_status)
if job_status in self.jobs_error_status:
raise AirflowException(
'Kylin job {} status {} is error '.format(
job_id, job_status))
if self.do_xcom_push:
return rsp_data