mtakaki · mtakaki · Mar 31, 2019 · Mar 19, 2019 · Mar 19, 2019 · Mar 21, 2019
diff --git a/README.md b/README.md
@@ -19,6 +19,8 @@ This project is available at PyPI: [https://pypi.python.org/pypi/cachet-url-moni
 endpoint:
   url: http://www.google.com
   method: GET
+  header:
+    SOME-HEADER: SOME-VALUE
   timeout: 1 # seconds
   expectation:
     - type: HTTP_STATUS
@@ -44,6 +46,7 @@ frequency: 30
 - **endpoint**, the configuration about the URL that will be monitored.
     - **url**, the URL that is going to be monitored.
     - **method**, the HTTP method that will be used by the monitor.
+    - **header**, client header passed to the request. Remove if you do not want to pass a header.
     - **timeout**, how long we'll wait to consider the request failed. The unit of it is seconds.
     - **expectation**, the list of expectations set for the URL.
         - **HTTP_STATUS**, we will verify if the response status code falls into the expected range. Please keep in mind the range is inclusive on the first number and exclusive on the second number. If just one value is specified, it will default to only the given value, for example `200` will be converted to `200-201`. 

diff --git a/cachet_url_monitor/configuration.py b/cachet_url_monitor/configuration.py
@@ -97,6 +97,7 @@ def __init__(self, config_file):
         self.endpoint_url = os.environ.get('ENDPOINT_URL') or self.data['endpoint']['url']
         self.endpoint_url = normalize_url(self.endpoint_url)
         self.endpoint_timeout = os.environ.get('ENDPOINT_TIMEOUT') or self.data['endpoint'].get('timeout') or 1
+        self.endpoint_header = self.data['endpoint'].get('header') or None
         self.allowed_fails = os.environ.get('ALLOWED_FAILS') or self.data['endpoint'].get('allowed_fails') or 0
 
         self.api_url = os.environ.get('CACHET_API_URL') or self.data['cachet']['api_url']
@@ -172,7 +173,10 @@ def evaluate(self):
         according to the expectation results.
         """
         try:
-            self.request = requests.request(self.endpoint_method, self.endpoint_url, timeout=self.endpoint_timeout)
+            if self.endpoint_header is not None:
+              self.request = requests.request(self.endpoint_method, self.endpoint_url, timeout=self.endpoint_timeout, headers=self.endpoint_header)
+            else:
+              self.request = requests.request(self.endpoint_method, self.endpoint_url, timeout=self.endpoint_timeout)
             self.current_timestamp = int(time.time())
         except requests.ConnectionError:
             self.message = 'The URL is unreachable: %s %s' % (self.endpoint_method, self.endpoint_url)

diff --git a/config.yml b/config.yml
@@ -1,6 +1,8 @@
 endpoint:
   url: http://localhost:8080/swagger
   method: GET
+  header:
+    SOME-HEADER: SOME-VALUE
   timeout: 0.01
   expectation:
     - type: HTTP_STATUS

diff --git a/tests/test_configuration.py b/tests/test_configuration.py
@@ -42,12 +42,31 @@ def test_init(self):
         self.assertDictEqual(self.configuration.headers, {'X-Cachet-Token': 'token2'}, 'Header was not set correctly')
         self.assertEquals(self.configuration.api_url, 'https://demo.cachethq.io/api/v1',
                           'Cachet API URL was set incorrectly')
+        self.assertDictEqual(self.configuration.endpoint_header, {'SOME-HEADER': 'SOME-VALUE'}, 'Header is incorrect')
 
     def test_evaluate(self):
         def total_seconds():
             return 0.1
 
-        def request(method, url, timeout=None):
+        def request(method, url, headers, timeout=None):
+            response = mock.Mock()
+            response.status_code = 200
+            response.elapsed = mock.Mock()
+            response.elapsed.total_seconds = total_seconds
+            response.text = '<body>'
+            return response
+
+        sys.modules['requests'].request = request
+        self.configuration.evaluate()
+
+        self.assertEquals(self.configuration.status, cachet_url_monitor.status.COMPONENT_STATUS_OPERATIONAL,
+                          'Component status set incorrectly')
+
+    def test_evaluate_without_header(self):
+        def total_seconds():
+            return 0.1
+
+        def request(method, url, headers=None, timeout=None):
             response = mock.Mock()
             response.status_code = 200
             response.elapsed = mock.Mock()
@@ -65,7 +84,7 @@ def test_evaluate_with_failure(self):
         def total_seconds():
             return 0.1
 
-        def request(method, url, timeout=None):
+        def request(method, url, headers, timeout=None):
             response = mock.Mock()
             # We are expecting a 200 response, so this will fail the expectation.
             response.status_code = 400
@@ -81,7 +100,7 @@ def request(method, url, timeout=None):
                           'Component status set incorrectly')
 
     def test_evaluate_with_timeout(self):
-        def request(method, url, timeout=None):
+        def request(method, url, headers, timeout=None):
             self.assertEquals(method, 'GET', 'Incorrect HTTP method')
             self.assertEquals(url, 'http://localhost:8080/swagger', 'Monitored URL is incorrect')
             self.assertEquals(timeout, 0.010)
@@ -96,7 +115,7 @@ def request(method, url, timeout=None):
         self.mock_logger.warning.assert_called_with('Request timed out')
 
     def test_evaluate_with_connection_error(self):
-        def request(method, url, timeout=None):
+        def request(method, url, headers, timeout=None):
             self.assertEquals(method, 'GET', 'Incorrect HTTP method')
             self.assertEquals(url, 'http://localhost:8080/swagger', 'Monitored URL is incorrect')
             self.assertEquals(timeout, 0.010)
@@ -111,7 +130,7 @@ def request(method, url, timeout=None):
         self.mock_logger.warning.assert_called_with('The URL is unreachable: GET http://localhost:8080/swagger')
 
     def test_evaluate_with_http_error(self):
-        def request(method, url, timeout=None):
+        def request(method, url, headers, timeout=None):
             self.assertEquals(method, 'GET', 'Incorrect HTTP method')
             self.assertEquals(url, 'http://localhost:8080/swagger', 'Monitored URL is incorrect')
             self.assertEquals(timeout, 0.010)