def _create_invalid_xcom_entries(self, execution_date, session=None): """ Invalid XCom entries to test join query """ for i in [1, 2]: XCom.set( key=f'invalid-xcom-key-{i}', value="TEST", execution_date=execution_date, task_id="invalid_task", dag_id="invalid_dag", ) dag = DagModel(dag_id="invalid_dag") session.add(dag) dagrun = DR( dag_id="invalid_dag", run_id="invalid_run_id", execution_date=execution_date + timedelta(days=1), start_date=execution_date, run_type=DagRunType.MANUAL, ) session.add(dagrun) dagrun = DR( dag_id="invalid_dag_1", run_id="invalid_run_id", execution_date=execution_date, start_date=execution_date, run_type=DagRunType.MANUAL, ) session.commit()
def test_should_response_200(self, session): dag_id = 'test-dag-id' task_id = 'test-task-id' execution_date = '2005-04-02T00:00:00+00:00' xcom_key = 'test-xcom-key' execution_date_parsed = parse_execution_date(execution_date) xcom_model = XCom(key=xcom_key, execution_date=execution_date_parsed, task_id=task_id, dag_id=dag_id, timestamp=execution_date_parsed) dag_run_id = DR.generate_run_id(DagRunType.MANUAL, execution_date_parsed) dagrun = DR(dag_id=dag_id, run_id=dag_run_id, execution_date=execution_date_parsed, start_date=execution_date_parsed, run_type=DagRunType.MANUAL.value) session.add(xcom_model) session.add(dagrun) session.commit() response = self.client.get( f"/api/v1/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/xcomEntries/{xcom_key}" ) self.assertEqual(200, response.status_code) self.assertEqual( response.json, { 'dag_id': dag_id, 'execution_date': execution_date, 'key': xcom_key, 'task_id': task_id, 'timestamp': execution_date })
def test_handle_limit_offset(self, query_params, expected_xcom_ids, session): url = "/api/v1/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/xcomEntries?{query_params}" url = url.format(dag_id=self.dag_id, dag_run_id=self.dag_run_id, task_id=self.task_id, query_params=query_params) dagrun = DR( dag_id=self.dag_id, run_id=self.dag_run_id, execution_date=self.execution_date_parsed, start_date=self.execution_date_parsed, run_type=DagRunType.MANUAL.value, ) xcom_models = self._create_xcoms(10) session.add_all(xcom_models) session.add(dagrun) session.commit() response = self.client.get(url, environ_overrides={'REMOTE_USER': "******"}) assert response.status_code == 200 self.assertEqual(response.json["total_entries"], 10) conn_ids = [ conn["key"] for conn in response.json["xcom_entries"] if conn ] self.assertEqual(conn_ids, expected_xcom_ids)
def _create_xcom_entries(self, dag_id, dag_run_id, execution_date, task_id, session=None): for i in [1, 2]: XCom.set( key=f'test-xcom-key-{i}', value="TEST", execution_date=execution_date, task_id=task_id, dag_id=dag_id, ) dag = DagModel(dag_id=dag_id) session.add(dag) dagrun = DR( dag_id=dag_id, run_id=dag_run_id, execution_date=execution_date, start_date=execution_date, run_type=DagRunType.MANUAL.value, ) session.add(dagrun)
def _create_xcom_entry(self, dag_id, dag_run_id, execution_date, task_id, xcom_key, session=None): XCom.set(key=xcom_key, value="TEST_VALUE", execution_date=execution_date, task_id=task_id, dag_id=dag_id,) dagrun = DR(dag_id=dag_id, run_id=dag_run_id, execution_date=execution_date, start_date=execution_date, run_type=DagRunType.MANUAL.value) session.add(dagrun)