From 0d1fd50a0c706c810357c745dd9d5c9b502ae7a1 Mon Sep 17 00:00:00 2001 From: Adam Welsh Date: Mon, 6 Aug 2018 16:16:48 +0100 Subject: [PATCH 1/2] [AIRFLOW-2860] Update tests for druid hook - Also assert that timeout is >= 1 --- airflow/hooks/druid_hook.py | 5 ++++- tests/hooks/test_druid_hook.py | 4 ++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/airflow/hooks/druid_hook.py b/airflow/hooks/druid_hook.py index e5fb6a7a26d91..4f1af46a2d82d 100644 --- a/airflow/hooks/druid_hook.py +++ b/airflow/hooks/druid_hook.py @@ -37,7 +37,8 @@ class DruidHook(BaseHook): which accepts index jobs :type druid_ingest_conn_id: string :param timeout: The interval between polling - the Druid job for the status of the ingestion job + the Druid job for the status of the ingestion job. + Must be greater than or equal to 1 :type timeout: int :param max_ingestion_time: The maximum ingestion time before assuming the job failed :type max_ingestion_time: int @@ -53,6 +54,8 @@ def __init__( self.max_ingestion_time = max_ingestion_time self.header = {'content-type': 'application/json'} + assert self.timeout >= 1 + def get_conn_url(self): conn = self.get_connection(self.druid_ingest_conn_id) host = conn.host diff --git a/tests/hooks/test_druid_hook.py b/tests/hooks/test_druid_hook.py index 6fd7b3cc76552..4243343b88227 100644 --- a/tests/hooks/test_druid_hook.py +++ b/tests/hooks/test_druid_hook.py @@ -98,7 +98,7 @@ def test_submit_unknown_response(self, m): @requests_mock.mock() def test_submit_timeout(self, m): - self.db_hook.timeout = 0 + self.db_hook.timeout = 1 self.db_hook.max_ingestion_time = 5 task_post = m.post( 'http://druid-overlord:8081/druid/indexer/v1/task', @@ -131,7 +131,7 @@ def test_get_conn_url(self, mock_get_connection): get_conn_value.port = '1' get_conn_value.extra_dejson = {'endpoint': 'ingest'} mock_get_connection.return_value = get_conn_value - hook = DruidHook(timeout=0, max_ingestion_time=5) + hook = DruidHook(timeout=1, max_ingestion_time=5) self.assertEquals(hook.get_conn_url(), 'https://test_host:1/ingest') From 1a8abf183738cbe3e2f073dd7909c34fcddce85c Mon Sep 17 00:00:00 2001 From: Adam Welsh Date: Tue, 7 Aug 2018 09:26:14 +0100 Subject: [PATCH 2/2] [AIRFLOW-2860] Raise ValueError if timeout < 1 in druid hook --- airflow/hooks/druid_hook.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/airflow/hooks/druid_hook.py b/airflow/hooks/druid_hook.py index 4f1af46a2d82d..fbd74e8e47d39 100644 --- a/airflow/hooks/druid_hook.py +++ b/airflow/hooks/druid_hook.py @@ -54,7 +54,8 @@ def __init__( self.max_ingestion_time = max_ingestion_time self.header = {'content-type': 'application/json'} - assert self.timeout >= 1 + if self.timeout < 1: + raise ValueError("Druid timeout should be equal or greater than 1") def get_conn_url(self): conn = self.get_connection(self.druid_ingest_conn_id)