def test_bigquery_operator_extra_link_when_single_query( self, mock_hook, session): bigquery_task = BigQueryExecuteQueryOperator( task_id=TASK_ID, sql='SELECT * FROM test_table', dag=self.dag, ) self.dag.clear() session.query(XCom).delete() ti = TaskInstance( task=bigquery_task, execution_date=DEFAULT_DATE, ) job_id = '12345' ti.xcom_push(key='job_id', value=job_id) self.assertEqual( 'https://console.cloud.google.com/bigquery?j={job_id}'.format( job_id=job_id), bigquery_task.get_extra_links(DEFAULT_DATE, BigQueryConsoleLink.name), ) self.assertEqual( '', bigquery_task.get_extra_links(datetime(2019, 1, 1), BigQueryConsoleLink.name), )
def test_bigquery_operator_extra_link_when_multiple_query( self, mock_hook, session): bigquery_task = BigQueryExecuteQueryOperator( task_id=TASK_ID, sql=['SELECT * FROM test_table', 'SELECT * FROM test_table2'], dag=self.dag, ) self.dag.clear() session.query(XCom).delete() ti = TaskInstance( task=bigquery_task, execution_date=DEFAULT_DATE, ) job_id = ['123', '45'] ti.xcom_push(key='job_id', value=job_id) self.assertEqual({'BigQuery Console #1', 'BigQuery Console #2'}, bigquery_task.operator_extra_link_dict.keys()) self.assertEqual( 'https://console.cloud.google.com/bigquery?j=123', bigquery_task.get_extra_links(DEFAULT_DATE, 'BigQuery Console #1'), ) self.assertEqual( 'https://console.cloud.google.com/bigquery?j=45', bigquery_task.get_extra_links(DEFAULT_DATE, 'BigQuery Console #2'), )
def test_bigquery_operator_extra_link_when_missing_job_id(self, mock_hook, session): bigquery_task = BigQueryExecuteQueryOperator( task_id=TASK_ID, sql='SELECT * FROM test_table', dag=self.dag, ) self.dag.clear() session.query(XCom).delete() self.assertEqual( '', bigquery_task.get_extra_links(DEFAULT_DATE, BigQueryConsoleLink.name), )