From 2df4f977bce1d31182cf57f62c0456f901dc8e0c Mon Sep 17 00:00:00 2001
From: Doug Daniels <doug@datadoghq.com>
Date: Wed, 8 Apr 2015 18:06:41 -0400
Subject: [PATCH] Sniff DATADOG_HOST environment variable to determine which
 API host to use

---
 datadog/__init__.py        |  6 ++++--
 tests/unit/api/test_api.py | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/datadog/__init__.py b/datadog/__init__.py
index 0d6fbdf65..8b3642468 100644
--- a/datadog/__init__.py
+++ b/datadog/__init__.py
@@ -8,6 +8,7 @@
 * datadog.dogshell: a command-line tool, wrapping datadog.api, to interact with Datadog REST API.
 """
 from pkg_resources import get_distribution, DistributionNotFound
+import os
 import os.path
 
 from datadog import api
@@ -30,7 +31,7 @@
     __version__ = _dist.version
 
 
-def initialize(api_key=None, app_key=None, host_name=None, api_host="https://app.datadoghq.com",
+def initialize(api_key=None, app_key=None, host_name=None, api_host=None,
                proxies=None, statsd_host=None, statsd_port=None):
     """
     Initialize and configure Datadog.api and Datadog.statsd modules
@@ -57,7 +58,8 @@ def initialize(api_key=None, app_key=None, host_name=None, api_host="https://app
     api._api_key = api_key
     api._application_key = app_key
     api._host_name = host_name if host_name is not None else get_hostname()
-    api._api_host = api_host
+    api._api_host = api_host if api_host is not None else \
+        os.environ.get('DATADOG_HOST', 'https://app.datadoghq.com')
     api._proxies = proxies
 
     # Given statsd_host and statsd_port, overrides statsd instance
diff --git a/tests/unit/api/test_api.py b/tests/unit/api/test_api.py
index ab217ae08..664649967 100644
--- a/tests/unit/api/test_api.py
+++ b/tests/unit/api/test_api.py
@@ -1,8 +1,10 @@
 # stdlib
+from functools import wraps
 import os
 import tempfile
 import time
 
+
 # 3p
 import mock
 from nose.tools import assert_raises, assert_true, assert_false
@@ -26,6 +28,23 @@
     HOST_NAME,
     FAKE_PROXY)
 
+def preserve_environ_datadog_host(func):
+    """
+    Decorator to preserve the original environment value of DATADOG_HOST.
+    """
+    @wraps(func)
+    def wrapper(*args, **kwds):
+        environ_api_host = os.environ.get('DATADOG_HOST')
+        try:
+            return func(*args, **kwds)
+        finally:
+            # restore the original environ value
+            if environ_api_host:
+                os.environ['DATADOG_HOST'] = environ_api_host
+            elif os.environ.get('DATADOG_HOST'):
+                del os.environ['DATADOG_HOST']
+
+    return wrapper
 
 class TestInitialization(DatadogAPINoInitialization):
 
@@ -90,6 +109,25 @@ def test_request_parameters(self):
         assert 'headers' in options
         assert options['headers'] == {'Content-Type': 'application/json'}
 
+    @preserve_environ_datadog_host
+    def test_api_host_from_env(self):
+        os.environ['DATADOG_HOST'] = 'http://localhost'
+        initialize()
+        self.assertEquals(api._api_host, 'http://localhost')
+
+    @preserve_environ_datadog_host
+    def test_api_host_default(self):
+        if os.environ.get('DATADOG_HOST'):
+            del os.environ['DATADOG_HOST']
+        initialize()
+        self.assertEquals(api._api_host, 'https://app.datadoghq.com')
+
+    @preserve_environ_datadog_host
+    def test_api_host_param(self):
+        if os.environ.get('DATADOG_HOST'):
+            del os.environ['DATADOG_HOST']
+        initialize(api_host='http://localhost')
+        self.assertEquals(api._api_host, 'http://localhost')
 
 class TestResources(DatadogAPIWithInitialization):