From e0420b8f9afe7b1c0e6dfd374b1f0c458a1f8eb9 Mon Sep 17 00:00:00 2001
From: Steven Yuen <steven.yuen@datadoghq.com>
Date: Thu, 24 Feb 2022 14:27:48 -0500
Subject: [PATCH] Adjust endpoint detection logic for warn log (#11567)

* change logic for endpoint detection

* check additional endpoints

* created helper function and wrote unit test for helper function

* rename function
---
 nginx/datadog_checks/nginx/nginx.py |  7 +++++-
 nginx/tests/test_unit.py            | 38 +++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/nginx/datadog_checks/nginx/nginx.py b/nginx/datadog_checks/nginx/nginx.py
index ea9d510aaa08f..d3dec2b5be041 100644
--- a/nginx/datadog_checks/nginx/nginx.py
+++ b/nginx/datadog_checks/nginx/nginx.py
@@ -281,7 +281,8 @@ def _get_plus_api_data(self, endpoint, nest):
             r = self._perform_request(url)
             payload = self._nest_payload(nest, r.json())
         except Exception as e:
-            if not self.only_query_enabled_endpoints and endpoint in PLUS_API_STREAM_ENDPOINTS.values():
+            plus_endpoints = self.list_endpoints(PLUS_API_STREAM_ENDPOINTS)
+            if not self.only_query_enabled_endpoints and endpoint in plus_endpoints:
                 self.log.warning(
                     "Error querying %s metrics at %s: %s. Stream may not be initialized, "
                     "you can avoid this error by enabling `only_query_enabled_endpoints` option.",
@@ -416,3 +417,7 @@ def _normalize_tags_type(self, tags, device_name=None, metric_name=None):
                     continue
                 normalized_tags.extend(list({tag, self.degeneralise_tag(tag)}))
         return normalized_tags
+
+    def list_endpoints(self, api_dict_list):
+        endpoints = [endpoint for api_dict in list(api_dict_list.values()) for endpoint in api_dict.keys()]
+        return endpoints
diff --git a/nginx/tests/test_unit.py b/nginx/tests/test_unit.py
index 186933667a196..427bd17522be3 100644
--- a/nginx/tests/test_unit.py
+++ b/nginx/tests/test_unit.py
@@ -247,3 +247,41 @@ def mock_get_return(*args, **kwargs):
         endpoints = check._get_enabled_endpoints()
         expected_endpoints = [('nginx', []), ('http/requests', ['requests'])]
         assert sorted(expected_endpoints) == sorted(list(endpoints))
+
+
+@pytest.mark.parametrize(
+    'test_input, expected_output',
+    [
+        (
+            {
+                '1': {
+                    "stream/server_zones": ["stream", "server_zones"],
+                    "stream/upstreams": ["stream", "upstreams"],
+                },
+                '3': {
+                    "stream/zone_sync": ["stream", "zone_sync"],
+                },
+                '6': {
+                    "stream/limit_conns": ["stream", "limit_conns"],
+                },
+            },
+            ['stream/server_zones', 'stream/upstreams', 'stream/zone_sync', 'stream/limit_conns'],
+        ),
+        (
+            {
+                'foo': {"biz": ["stream1"], "buz": ["stream1", "stream2"], "bes": "stream3"},
+                "baz": {"bux": "zone_sync", "bus": "zone_sync"},
+                "bar": {
+                    "bis": ["stream1", "stream2", "stream3"],
+                },
+            },
+            ['biz', 'buz', 'bes', 'bux', 'bus', 'bis'],
+        ),
+    ],
+)
+def test_list_endpoints(instance, test_input, expected_output):
+    nginx = Nginx('nginx', {}, [instance])
+    # Python 2 seems to have some different order of processing the keys.
+    # Sorting the arrays before comparison to account for this.
+    sorted_test_output = nginx.list_endpoints(test_input).sort()
+    assert eval(str(sorted_test_output)) == expected_output.sort()