def test_bigquery_operator_extra_link_when_single_query(
self, mock_hook, session):
bigquery_task = BigQueryOperator(
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 = BigQueryOperator(
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 = BigQueryOperator(
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),
)