def setUp(self, mock_get_connection, mock_initialize):
mock_get_connection.return_value = Connection(
extra=json.dumps({
'app_key': APP_KEY,
'api_key': API_KEY,
}))
self.hook = DatadogHook()
def test_api_key_required(self, mock_get_connection, mock_initialize):
mock_get_connection.return_value = Connection()
with pytest.raises(AirflowException) as ctx:
DatadogHook()
assert str(
ctx.value
) == 'api_key must be specified in the Datadog connection details'
def poke(self, context: 'Context') -> bool:
# This instantiates the hook, but doesn't need it further,
# because the API authenticates globally (unfortunately),
# but for airflow this shouldn't matter too much, because each
# task instance runs in its own process anyway.
DatadogHook(datadog_conn_id=self.datadog_conn_id)
response = api.Event.query(
start=self.from_seconds_ago,
end=self.up_to_seconds_from_now,
priority=self.priority,
sources=self.sources,
tags=self.tags,
)
if isinstance(response, dict) and response.get('status', 'ok') != 'ok':
self.log.error("Unexpected Datadog result: %s", response)
raise AirflowException("Datadog returned unexpected result")
if self.response_check:
# run content check on response
return self.response_check(response)
# If no check was inserted, assume any event that matched yields true.
return len(response) > 0
def test_api_key_required(self, mock_get_connection, mock_initialize):
mock_get_connection.return_value = Connection()
with self.assertRaises(AirflowException) as ctx:
DatadogHook()
self.assertEqual(
str(ctx.exception),
'api_key must be specified in the Datadog connection details')
class TestDatadogHook(unittest.TestCase):
@mock.patch('airflow.providers.datadog.hooks.datadog.initialize')
@mock.patch(
'airflow.providers.datadog.hooks.datadog.DatadogHook.get_connection')
def setUp(self, mock_get_connection, mock_initialize):
mock_get_connection.return_value = Connection(
extra=json.dumps({
'app_key': APP_KEY,
'api_key': API_KEY,
}))
self.hook = DatadogHook()
@mock.patch('airflow.providers.datadog.hooks.datadog.initialize')
@mock.patch(
'airflow.providers.datadog.hooks.datadog.DatadogHook.get_connection')
def test_api_key_required(self, mock_get_connection, mock_initialize):
mock_get_connection.return_value = Connection()
with pytest.raises(AirflowException) as ctx:
DatadogHook()
assert str(
ctx.value
) == 'api_key must be specified in the Datadog connection details'
def test_validate_response_valid(self):
try:
self.hook.validate_response({'status': 'ok'})
except AirflowException:
self.fail('Unexpected AirflowException raised')
def test_validate_response_invalid(self):
with pytest.raises(AirflowException):
self.hook.validate_response({'status': 'error'})
@mock.patch('airflow.providers.datadog.hooks.datadog.api.Metric.send')
def test_send_metric(self, mock_send):
mock_send.return_value = {'status': 'ok'}
self.hook.send_metric(
METRIC_NAME,
DATAPOINT,
tags=TAGS,
type_=TYPE,
interval=INTERVAL,
)
mock_send.assert_called_once_with(
metric=METRIC_NAME,
points=DATAPOINT,
host=self.hook.host,
tags=TAGS,
type=TYPE,
interval=INTERVAL,
)
@mock.patch('airflow.providers.datadog.hooks.datadog.api.Metric.query')
@mock.patch('airflow.providers.datadog.hooks.datadog.time.time')
def test_query_metric(self, mock_time, mock_query):
now = 12345
mock_time.return_value = now
mock_query.return_value = {'status': 'ok'}
self.hook.query_metric('query', 60, 30)
mock_query.assert_called_once_with(
start=now - 60,
end=now - 30,
query='query',
)
@mock.patch('airflow.providers.datadog.hooks.datadog.api.Event.create')
def test_post_event(self, mock_create):
mock_create.return_value = {'status': 'ok'}
self.hook.post_event(
TITLE,
TEXT,
aggregation_key=AGGREGATION_KEY,
alert_type=ALERT_TYPE,
date_happened=DATE_HAPPENED,
handle=HANDLE,
priority=PRIORITY,
related_event_id=RELATED_EVENT_ID,
tags=TAGS,
device_name=DEVICE_NAME,
)
mock_create.assert_called_once_with(
title=TITLE,
text=TEXT,
aggregation_key=AGGREGATION_KEY,
alert_type=ALERT_TYPE,
date_happened=DATE_HAPPENED,
handle=HANDLE,
priority=PRIORITY,
related_event_id=RELATED_EVENT_ID,
tags=TAGS,
host=self.hook.host,
device_name=DEVICE_NAME,
source_type_name=self.hook.source_type_name,
)
