dsquery: Adjust software tests

Author: amotl

grafana_client/elements/datasource.py

@@ -16,6 +16,7 @@
 logger = logging.getLogger(__name__)
 
 VERSION_9 = LooseVersion("9")
+VERSION_8 = LooseVersion("8")
 VERSION_7 = LooseVersion("7")
 VERBOSE = False
 
@@ -394,6 +395,7 @@ def smartquery(
                     request["data"]["to"],
                     request["data"]["step"],
                 )
+
         # For all others, use the generic data source communication endpoint.
         elif access_type in ["server", "proxy"]:
             url = "/ds/query"
@@ -438,6 +440,7 @@ def health_check(self, datasource: Union[DatasourceIdentifier, Dict]) -> Datasou
         expression = get_healthcheck_expression(datasource_type, datasource_dialect)
 
         start = time.time()
+        message = "Unknown error"
         try:
             response = self.smartquery(datasource, expression)
             response_display = response
@@ -487,15 +490,16 @@ def health_check(self, datasource: Union[DatasourceIdentifier, Dict]) -> Datasou
                     message = f"Invalid response. {reason}"
 
             elif datasource_type == "loki":
-                if self.api.version:
-                    if LooseVersion(self.api.version) == VERSION_7:
-                        if "status" in response and response["status"] == "success":
-                            message = "Success"
-                            success = True
-                    elif "results" in response and "test" in response["results"]:
+                if self.api.version and VERSION_7 <= LooseVersion(self.api.version) < VERSION_8:
+                    if "status" in response and response["status"] == "success":
                         message = "Success"
                         success = True
-                else:
+                    else:
+                        message = response.get("message", "Unknown error")
+                elif "results" in response and "test" in response["results"]:
+                    message = "Success"
+                    success = True
+                elif "message" in response:
                     message = response["message"]
 
             # With OpenTSDB, a 200 OK response with empty body is just fine.

grafana_client/knowledge.py

@@ -115,7 +115,7 @@ def datasource_factory(datasource: DatasourceModel) -> DatasourceModel:
     return datasource
 
 
-def query_factory(datasource, model: Optional[dict]) -> Union[Dict, str]:
+def query_factory(datasource, model: Optional[dict] = None, expression: Optional[str] = None) -> Union[Dict, str]:
     """
     Create payload suitable for running a query against a Grafana data source.
 
@@ -124,6 +124,11 @@ def query_factory(datasource, model: Optional[dict]) -> Union[Dict, str]:
 
     TODO: Complete the list for all popular databases.
     """
+
+    model = model or {}
+    if "query" not in model and expression:
+        model["query"] = expression
+
     request = {
         "method": "POST",
         "data": None,

test/elements/test_datasource_base.py

@@ -1,5 +1,6 @@
 import unittest
 from test.elements.test_datasource_fixtures import (
+    DATAFRAME_RESPONSE_HEALTH_ELASTICSEARCH_VALID,
     DATAFRAME_RESPONSE_HEALTH_PROMETHEUS,
     ELASTICSEARCH_DATASOURCE,
     INFLUXDB1_DATASOURCE,
@@ -271,7 +272,29 @@ def test_series(self, m):
         self.assertEqual(len(result["data"]["result"][0]["values"]), 6)
 
     @requests_mock.Mocker()
-    def test_query_with_datasource_prometheus(self, m):
+    def test_query_with_datasource_prometheus_grafana7(self, m):
+        # Mock the version inquiry request, because `smartquery` needs
+        # it, as Prometheus responses differ between versions.
+        m.get(
+            "http://localhost/api/health",
+            json={"commit": "unknown", "database": "ok", "version": "7.0.1"},
+        )
+        m.post(
+            "http://localhost/api/ds/query",
+            json=DATAFRAME_RESPONSE_HEALTH_PROMETHEUS,
+        )
+        datasource = PROMETHEUS_DATASOURCE.copy()
+        response = self.grafana.datasource.smartquery(datasource, "1+1")
+        self.assertEqual(response, DATAFRAME_RESPONSE_HEALTH_PROMETHEUS)
+
+    @requests_mock.Mocker()
+    def test_query_with_datasource_prometheus_grafana9(self, m):
+        # Mock the version inquiry request, because `smartquery` needs
+        # it, as Prometheus responses differ between versions.
+        m.get(
+            "http://localhost/api/health",
+            json={"commit": "14e988bd22", "database": "ok", "version": "9.0.1"},
+        )
         m.post(
             "http://localhost/api/ds/query",
             json=DATAFRAME_RESPONSE_HEALTH_PROMETHEUS,
@@ -298,12 +321,22 @@ def test_query_with_datasource_elasticsearch(self, m):
             "http://localhost/api/datasources/proxy/44/bazqux/_mapping",
             json={},
         )
+        m.post(
+            "http://localhost/api/ds/query",
+            json=DATAFRAME_RESPONSE_HEALTH_ELASTICSEARCH_VALID,
+        )
         datasource = ELASTICSEARCH_DATASOURCE.copy()
         _ = self.grafana.datasource.smartquery(datasource, "url:///datasources/proxy/44/bazqux/_mapping")
-        # TODO: No response payload yet.
+        # TODO: Response payload not reflected and validated yet.
 
     @requests_mock.Mocker()
     def test_query_with_datasource_identifier(self, m):
+        # Mock the version inquiry request, because `smartquery` needs
+        # it, as Prometheus responses differ between versions.
+        m.get(
+            "http://localhost/api/health",
+            json={"commit": "14e988bd22", "database": "ok", "version": "9.0.1"},
+        )
         m.get(
             "http://localhost/api/datasources/uid/h8KkCLt7z",
             json=PROMETHEUS_DATASOURCE,
@@ -315,7 +348,14 @@ def test_query_with_datasource_identifier(self, m):
         response = self.grafana.datasource.smartquery(DatasourceIdentifier(uid="h8KkCLt7z"), "1+1")
         self.assertEqual(response, DATAFRAME_RESPONSE_HEALTH_PROMETHEUS)
 
-    def test_query_unknown_access_type_failure(self):
+    @requests_mock.Mocker()
+    def test_query_unknown_access_type_failure(self, m):
+        # Mock the version inquiry request, because `smartquery` needs
+        # it, as Prometheus responses differ between versions.
+        m.get(
+            "http://localhost/api/health",
+            json={"commit": "14e988bd22", "database": "ok", "version": "9.0.1"},
+        )
         datasource = PROMETHEUS_DATASOURCE.copy()
         datasource["access"] = "__UNKNOWN__"
         self.assertRaises(NotImplementedError, lambda: self.grafana.datasource.smartquery(datasource, expression="1+1"))

test/elements/test_datasource_fixtures.py

@@ -205,6 +205,9 @@
     "results": {"test": {"frames": [{"schema": {"meta": {"executedQueryString": "SELECT 1"}}}]}}
 }
 
+# TODO: Reflect a real Grafana response for Elasticsearch here.
+DATAFRAME_RESPONSE_HEALTH_ELASTICSEARCH_VALID = {"bazqux": {}}
+
 DATAFRAME_RESPONSE_HEALTH_PROMETHEUS = {
     "results": {
         "test": {