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 self.assertRaises(AirflowException) as ctx:
DatadogHook()
self.assertEqual(
str(ctx.exception),
'api_key must be specified in the Datadog connection details')
def poke(self, context):
# 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
class TestDatadogHook(unittest.TestCase):
@mock.patch('airflow.contrib.hooks.datadog_hook.initialize')
@mock.patch('airflow.contrib.hooks.datadog_hook.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.contrib.hooks.datadog_hook.initialize')
@mock.patch('airflow.contrib.hooks.datadog_hook.DatadogHook.get_connection'
)
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')
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 self.assertRaises(AirflowException):
self.hook.validate_response({'status': 'error'})
@mock.patch('airflow.contrib.hooks.datadog_hook.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_with(
metric=METRIC_NAME,
points=DATAPOINT,
host=self.hook.host,
tags=TAGS,
type=TYPE,
interval=INTERVAL,
)
@mock.patch('airflow.contrib.hooks.datadog_hook.api.Metric.query')
@mock.patch('airflow.contrib.hooks.datadog_hook.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_with(
start=now - 60,
end=now - 30,
query='query',
)
@mock.patch('airflow.contrib.hooks.datadog_hook.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_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,
)
from airflow.contrib.hooks.datadog_hook import DatadogHook
hook = DatadogHook()
title = "Test airflow event"
text = "This is a test event from the airflow hook"
alert_type = "error"
tags = ["airflow_test"]
#test event to the org event stream
#hook.post_event(title, text, aggregation_key=None, alert_type=alert_type, date_happened=None, \
# handle=None, priority=None, related_event_id=None, tags=tags, device_name=None)
metric_name = "airflow_hook_test_metric"
datapoint = 1
tags = ["airflow_test"]
metric_type = "count"
#send a metric from the hook
#hook.send_metric(metric_name, datapoint, tags=tags, type_=metric_type, interval=None)
query = "avg:airflow_hook_test_metric{*}.as_count()"
from_seconds_ago = 3600
to_seconds_ago = 0
response = {}
#query a metric sent via the airflow datadog hook
#response = hook.query_metric(query, from_seconds_ago, to_seconds_ago)
#print(response)