|
1 | 1 | # stdlib |
| 2 | +from functools import wraps |
2 | 3 | import os |
3 | 4 | import tempfile |
4 | 5 | import time |
5 | 6 |
|
| 7 | + |
6 | 8 | # 3p |
7 | 9 | import mock |
8 | 10 | from nose.tools import assert_raises, assert_true, assert_false |
|
26 | 28 | HOST_NAME, |
27 | 29 | FAKE_PROXY) |
28 | 30 |
|
| 31 | +def preserve_environ_datadog_host(func): |
| 32 | + """ |
| 33 | + Decorator to preserve the original environment value of DATADOG_HOST. |
| 34 | + """ |
| 35 | + @wraps(func) |
| 36 | + def wrapper(*args, **kwds): |
| 37 | + environ_api_host = os.environ.get('DATADOG_HOST') |
| 38 | + try: |
| 39 | + return func(*args, **kwds) |
| 40 | + finally: |
| 41 | + # restore the original environ value |
| 42 | + if environ_api_host: |
| 43 | + os.environ['DATADOG_HOST'] = environ_api_host |
| 44 | + elif os.environ.get('DATADOG_HOST'): |
| 45 | + del os.environ['DATADOG_HOST'] |
| 46 | + |
| 47 | + return wrapper |
29 | 48 |
|
30 | 49 | class TestInitialization(DatadogAPINoInitialization): |
31 | 50 |
|
@@ -90,6 +109,25 @@ def test_request_parameters(self): |
90 | 109 | assert 'headers' in options |
91 | 110 | assert options['headers'] == {'Content-Type': 'application/json'} |
92 | 111 |
|
| 112 | + @preserve_environ_datadog_host |
| 113 | + def test_api_host_from_env(self): |
| 114 | + os.environ['DATADOG_HOST'] = 'http://localhost' |
| 115 | + initialize() |
| 116 | + self.assertEquals(api._api_host, 'http://localhost') |
| 117 | + |
| 118 | + @preserve_environ_datadog_host |
| 119 | + def test_api_host_default(self): |
| 120 | + if os.environ.get('DATADOG_HOST'): |
| 121 | + del os.environ['DATADOG_HOST'] |
| 122 | + initialize() |
| 123 | + self.assertEquals(api._api_host, 'https://app.datadoghq.com') |
| 124 | + |
| 125 | + @preserve_environ_datadog_host |
| 126 | + def test_api_host_param(self): |
| 127 | + if os.environ.get('DATADOG_HOST'): |
| 128 | + del os.environ['DATADOG_HOST'] |
| 129 | + initialize(api_host='http://localhost') |
| 130 | + self.assertEquals(api._api_host, 'http://localhost') |
93 | 131 |
|
94 | 132 | class TestResources(DatadogAPIWithInitialization): |
95 | 133 |
|
|
0 commit comments