def test_should_pass_message_to_client(self, mock_client):
transport_type = mock.MagicMock()
stackdriver_task_handler = StackdriverTaskHandler(
gcp_conn_id=None,
transport=transport_type,
labels={"key": 'value'}
)
logger = logging.getLogger("logger")
logger.addHandler(stackdriver_task_handler)
logger.info("test-message")
stackdriver_task_handler.flush()
transport_type.assert_called_once_with(mock_client.return_value, 'airflow')
transport_type.return_value.send.assert_called_once_with(
mock.ANY, 'test-message', labels={"key": 'value'}, resource=Resource(type='global', labels={})
)
mock_client.assert_called_once()
class TestStackdriverLoggingHandlerTask(unittest.TestCase):
def setUp(self) -> None:
self.transport_mock = mock.MagicMock()
self.stackdriver_task_handler = StackdriverTaskHandler(
transport=self.transport_mock)
self.logger = logging.getLogger("logger")
date = datetime(2016, 1, 1)
self.dag = DAG('dag_for_testing_file_task_handler', start_date=date)
task = DummyOperator(task_id='task_for_testing_file_log_handler',
dag=self.dag)
self.ti = TaskInstance(task=task, execution_date=date)
self.ti.try_number = 1
self.ti.state = State.RUNNING
self.addCleanup(self.dag.clear)
@mock.patch(
'airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id'
)
@mock.patch('airflow.utils.log.stackdriver_task_handler.gcp_logging.Client'
)
def test_should_set_labels(self, mock_client,
mock_get_creds_and_project_id):
self.stackdriver_task_handler.set_context(self.ti)
self.logger.addHandler(self.stackdriver_task_handler)
self.logger.info("test-message")
self.stackdriver_task_handler.flush()
labels = {
'task_id': 'task_for_testing_file_log_handler',
'dag_id': 'dag_for_testing_file_task_handler',
'execution_date': '2016-01-01T00:00:00+00:00',
'try_number': '1'
}
resource = Resource(type='global', labels={})
self.transport_mock.return_value.send.assert_called_once_with(
mock.ANY, 'test-message', labels=labels, resource=resource)
@mock.patch(
'airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id'
)
@mock.patch('airflow.utils.log.stackdriver_task_handler.gcp_logging.Client'
)
def test_should_append_labels(self, mock_client,
mock_get_creds_and_project_id):
self.stackdriver_task_handler = StackdriverTaskHandler(
transport=self.transport_mock,
labels={"product.googleapis.com/task_id": "test-value"})
self.stackdriver_task_handler.set_context(self.ti)
self.logger.addHandler(self.stackdriver_task_handler)
self.logger.info("test-message")
self.stackdriver_task_handler.flush()
labels = {
'task_id': 'task_for_testing_file_log_handler',
'dag_id': 'dag_for_testing_file_task_handler',
'execution_date': '2016-01-01T00:00:00+00:00',
'try_number': '1',
'product.googleapis.com/task_id': 'test-value'
}
resource = Resource(type='global', labels={})
self.transport_mock.return_value.send.assert_called_once_with(
mock.ANY, 'test-message', labels=labels, resource=resource)
@mock.patch(
'airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id'
)
@mock.patch(
'airflow.utils.log.stackdriver_task_handler.gcp_logging.Client',
**{'return_value.project': 'asf-project'} # type: ignore
)
def test_should_read_logs_for_all_try(self, mock_client,
mock_get_creds_and_project_id):
mock_client.return_value.list_entries.return_value = _create_list_response(
["MSG1", "MSG2"], None)
logs, metadata = self.stackdriver_task_handler.read(self.ti)
mock_client.return_value.list_entries.assert_called_once_with(
filter_='resource.type="global"\n'
'logName="projects/asf-project/logs/airflow"\n'
'labels.task_id="task_for_testing_file_log_handler"\n'
'labels.dag_id="dag_for_testing_file_task_handler"\n'
'labels.execution_date="2016-01-01T00:00:00+00:00"',
page_token=None)
self.assertEqual(['MSG1\nMSG2'], logs)
self.assertEqual([{'end_of_log': True}], metadata)
@mock.patch(
'airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id'
)
@mock.patch( # type: ignore
'airflow.utils.log.stackdriver_task_handler.gcp_logging.Client',
**{'return_value.project': 'asf-project'} # type: ignore
)
def test_should_read_logs_for_task_with_quote(
self, mock_client, mock_get_creds_and_project_id):
mock_client.return_value.list_entries.return_value = _create_list_response(
["MSG1", "MSG2"], None)
self.ti.task_id = "K\"OT"
logs, metadata = self.stackdriver_task_handler.read(self.ti)
mock_client.return_value.list_entries.assert_called_once_with(
filter_='resource.type="global"\n'
'logName="projects/asf-project/logs/airflow"\n'
'labels.task_id="K\\"OT"\n'
'labels.dag_id="dag_for_testing_file_task_handler"\n'
'labels.execution_date="2016-01-01T00:00:00+00:00"',
page_token=None)
self.assertEqual(['MSG1\nMSG2'], logs)
self.assertEqual([{'end_of_log': True}], metadata)
@mock.patch(
'airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id'
)
@mock.patch(
'airflow.utils.log.stackdriver_task_handler.gcp_logging.Client',
**{'return_value.project': 'asf-project'} # type: ignore
)
def test_should_read_logs_for_single_try(self, mock_client,
mock_get_creds_and_project_id):
mock_client.return_value.list_entries.return_value = _create_list_response(
["MSG1", "MSG2"], None)
logs, metadata = self.stackdriver_task_handler.read(self.ti, 3)
mock_client.return_value.list_entries.assert_called_once_with(
filter_='resource.type="global"\n'
'logName="projects/asf-project/logs/airflow"\n'
'labels.task_id="task_for_testing_file_log_handler"\n'
'labels.dag_id="dag_for_testing_file_task_handler"\n'
'labels.execution_date="2016-01-01T00:00:00+00:00"\n'
'labels.try_number="3"',
page_token=None)
self.assertEqual(['MSG1\nMSG2'], logs)
self.assertEqual([{'end_of_log': True}], metadata)
@mock.patch(
'airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id'
)
@mock.patch('airflow.utils.log.stackdriver_task_handler.gcp_logging.Client'
)
def test_should_read_logs_with_pagination(self, mock_client,
mock_get_creds_and_project_id):
mock_client.return_value.list_entries.side_effect = [
_create_list_response(["MSG1", "MSG2"], "TOKEN1"),
_create_list_response(["MSG3", "MSG4"], None),
]
logs, metadata1 = self.stackdriver_task_handler.read(self.ti, 3)
mock_client.return_value.list_entries.assert_called_once_with(
filter_=mock.ANY, page_token=None)
self.assertEqual(['MSG1\nMSG2'], logs)
self.assertEqual([{
'end_of_log': False,
'next_page_token': 'TOKEN1'
}], metadata1)
mock_client.return_value.list_entries.return_value.next_page_token = None
logs, metadata2 = self.stackdriver_task_handler.read(
self.ti, 3, metadata1[0])
mock_client.return_value.list_entries.assert_called_with(
filter_=mock.ANY, page_token="TOKEN1")
self.assertEqual(['MSG3\nMSG4'], logs)
self.assertEqual([{'end_of_log': True}], metadata2)
@mock.patch(
'airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id'
)
@mock.patch('airflow.utils.log.stackdriver_task_handler.gcp_logging.Client'
)
def test_should_read_logs_with_download(self, mock_client,
mock_get_creds_and_project_id):
mock_client.return_value.list_entries.side_effect = [
_create_list_response(["MSG1", "MSG2"], "TOKEN1"),
_create_list_response(["MSG3", "MSG4"], None),
]
logs, metadata1 = self.stackdriver_task_handler.read(
self.ti, 3, {'download_logs': True})
self.assertEqual(['MSG1\nMSG2\nMSG3\nMSG4'], logs)
self.assertEqual([{'end_of_log': True}], metadata1)
@mock.patch(
'airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id'
)
@mock.patch(
'airflow.utils.log.stackdriver_task_handler.gcp_logging.Client',
**{'return_value.project': 'asf-project'} # type: ignore
)
def test_should_read_logs_with_custom_resources(
self, mock_client, mock_get_creds_and_project_id):
resource = Resource(
type="cloud_composer_environment",
labels={
"environment.name": 'test-instancce',
"location": 'europpe-west-3',
"project_id": "asf-project",
},
)
self.stackdriver_task_handler = StackdriverTaskHandler(
transport=self.transport_mock, resource=resource)
entry = mock.MagicMock(payload={"message": "TEXT"})
page = [entry, entry]
mock_client.return_value.list_entries.return_value.pages = (
n for n in [page])
mock_client.return_value.list_entries.return_value.next_page_token = None
logs, metadata = self.stackdriver_task_handler.read(self.ti)
mock_client.return_value.list_entries.assert_called_once_with(
filter_='resource.type="cloud_composer_environment"\n'
'logName="projects/asf-project/logs/airflow"\n'
'resource.labels."environment.name"="test-instancce"\n'
'resource.labels.location="europpe-west-3"\n'
'resource.labels.project_id="asf-project"\n'
'labels.task_id="task_for_testing_file_log_handler"\n'
'labels.dag_id="dag_for_testing_file_task_handler"\n'
'labels.execution_date="2016-01-01T00:00:00+00:00"',
page_token=None)
self.assertEqual(['TEXT\nTEXT'], logs)
self.assertEqual([{'end_of_log': True}], metadata)
@mock.patch(
'airflow.utils.log.stackdriver_task_handler.get_credentials_and_project_id'
)
@mock.patch('airflow.utils.log.stackdriver_task_handler.gcp_logging.Client'
)
def test_should_use_credentials(self, mock_client,
mock_get_creds_and_project_id):
stackdriver_task_handler = StackdriverTaskHandler(
gcp_key_path="KEY_PATH", )
client = stackdriver_task_handler._client
mock_get_creds_and_project_id.assert_called_once_with(
key_path='KEY_PATH',
scopes=frozenset({
'https://www.googleapis.com/auth/logging.write',
'https://www.googleapis.com/auth/logging.read'
}))
mock_client.assert_called_once_with(
credentials=mock_get_creds_and_project_id.return_value,
client_info=mock.ANY)
self.assertEqual(mock_client.return_value, client)