jprdonnelly/redash
1
0
Fork 0
mirror of https://github.com/getredash/redash.git synced 2025-12-19 17:37:19 -05:00
Code Issues Projects Releases Wiki Activity
Files
58bf96c2983f385db1053f14fc4ee3a0a16bd010
redash/tests/query_runner/test_prometheus.py
fabrei 58bf96c298 Adds ssl support for prometheus query runner. (#6657) 
* Adds ssl support for prometheus query runner.

- Adds possibilty to upload and use of ssl cert, key and ca file in redash ui

* Extends test cases for prometheus query runner.

- Adds secret attribute to configuration schema.

* Fixes wrong timestamps in different timezones in prometheus' testcases.

- Dynamically calculates timestamps in testcases to be robust in
  different timezones.
- Adds now datetime function to make it more testable.

* Fixes timestamp in prometheus' testcases which can be wrong depending on timezone.

---------

Co-authored-by: Masayuki Takahashi <masayuki038@gmail.com>
2023-12-17 21:58:16 +09:00

541 lines
20 KiB
Python
Raw Blame History

import time
from datetime import datetime
from unittest import TestCase
import mock
from redash.query_runner.prometheus import Prometheus, get_instant_rows, get_range_rows
from redash.utils import json_dumps
class TestPrometheus(TestCase):
def setUp(self):
self.instant_query_result = [
{
"metric": {"name": "example_metric_name", "foo_bar": "foo"},
"value": [1516937400.781, "7400_foo"],
},
{
"metric": {"name": "example_metric_name", "foo_bar": "bar"},
"value": [1516937400.781, "7400_bar"],
},
]
self.range_query_result = [
{
"metric": {"name": "example_metric_name", "foo_bar": "foo"},
"values": [[1516937400.781, "7400_foo"], [1516938000.781, "8000_foo"]],
},
{
"metric": {"name": "example_metric_name", "foo_bar": "bar"},
"values": [[1516937400.781, "7400_bar"], [1516938000.781, "8000_bar"]],
},
]
def test_get_instant_rows(self):
instant_rows = [
{
"name": "example_metric_name",
"foo_bar": "foo",
"timestamp": datetime.fromtimestamp(1516937400.781),
"value": "7400_foo",
},
{
"name": "example_metric_name",
"foo_bar": "bar",
"timestamp": datetime.fromtimestamp(1516937400.781),
"value": "7400_bar",
},
]
rows = get_instant_rows(self.instant_query_result)
self.assertEqual(instant_rows, rows)
def test_get_range_rows(self):
range_rows = [
{
"name": "example_metric_name",
"foo_bar": "foo",
"timestamp": datetime.fromtimestamp(1516937400.781),
"value": "7400_foo",
},
{
"name": "example_metric_name",
"foo_bar": "foo",
"timestamp": datetime.fromtimestamp(1516938000.781),
"value": "8000_foo",
},
{
"name": "example_metric_name",
"foo_bar": "bar",
"timestamp": datetime.fromtimestamp(1516937400.781),
"value": "7400_bar",
},
{
"name": "example_metric_name",
"foo_bar": "bar",
"timestamp": datetime.fromtimestamp(1516938000.781),
"value": "8000_bar",
},
]
rows = get_range_rows(self.range_query_result)
self.assertEqual(range_rows, rows)
@mock.patch("redash.query_runner.prometheus.datetime")
def test_get_datetime_now(self, datetime_mock: mock.MagicMock):
prometheus = Prometheus({"url": "url"})
datetime_mock.now.return_value = datetime(2023, 12, 12, 11, 00)
now = prometheus._get_datetime_now()
self.assertEqual(now, datetime(2023, 12, 12, 11, 00))
@mock.patch("redash.query_runner.prometheus.Prometheus._create_cert_file")
def test_get_prometheus_kwargs(self, create_cert_file_mock: mock.MagicMock):
# 1. case: without ssl attributes
prometheus = Prometheus({"url": "url"})
create_cert_file_mock.return_value = None
prometheus_kwargs = prometheus._get_prometheus_kwargs()
assert prometheus_kwargs == {
"verify": True,
"cert": (),
}
create_cert_file_mock.assert_has_calls(
[mock.call("ca_cert_File"), mock.call("cert_File"), mock.call("cert_key_File")]
)
create_cert_file_mock.reset_mock()
# 2. case: with ssl attributes
create_cert_file_return_dict = {
"ca_cert_File": "ca_cert_file.crt",
"cert_File": "cert_file.crt",
"cert_key_File": "cert_key_file.key",
}
create_cert_file_mock.side_effect = lambda key: create_cert_file_return_dict[key]
prometheus = Prometheus(
{
"url": "url",
"verify_ssl": True,
"ca_cert_File": "ca_cert_file",
"cert_File": "cert_file",
"cert_key_File": "cert_key_file",
}
)
prometheus_kwargs = prometheus._get_prometheus_kwargs()
assert prometheus_kwargs == {
"verify": "ca_cert_file.crt",
"cert": ("cert_file.crt", "cert_key_file.key"),
}
create_cert_file_mock.assert_has_calls(
[mock.call("ca_cert_File"), mock.call("cert_File"), mock.call("cert_key_File")]
)
@mock.patch("redash.query_runner.prometheus.NamedTemporaryFile")
def test_create_cert_file(self, named_temporary_file_mock: mock.MagicMock):
# 1. case: with none value
prometheus = Prometheus({"url": "url"})
context_manager_mock = named_temporary_file_mock().__enter__()
cert_file_name = prometheus._create_cert_file("key")
assert cert_file_name is None
context_manager_mock.write().assert_not_called()
named_temporary_file_mock.reset_mock()
# 2. case: with a valid key
prometheus = Prometheus({"url": "url", "key": "dmFsdWU="})
context_manager_mock = named_temporary_file_mock().__enter__()
context_manager_mock.name = "cert_file_name"
cert_file_name = prometheus._create_cert_file("key")
assert cert_file_name == "cert_file_name"
context_manager_mock.write.assert_called_once_with("value")
@mock.patch("redash.query_runner.prometheus.os")
def test_cleanup_cert_files(self, os_mock: mock.MagicMock):
# 1. case: no file found or verify is bool
prometheus = Prometheus(
{
"url": "url",
"verify_ssl": True,
"ca_cert_File": "ca_cert_file",
"cert_File": "cert_file",
"cert_key_File": "cert_key_file",
}
)
prometheus._cleanup_cert_files({"verify": True, "cert": ()})
os_mock.path.exists.assert_not_called()
os_mock.remove.assert_not_called()
# 2. case: files found and deleted
os_mock.path.exists.return_value = True
prometheus._cleanup_cert_files({"verify": "ca_cert_file", "cert": ("cert_file", "cert_key_file")})
os_mock.path.exists.assert_has_calls(
[mock.call("ca_cert_file"), mock.call("cert_file"), mock.call("cert_key_file")]
)
os_mock.remove.assert_has_calls(
[mock.call("ca_cert_file"), mock.call("cert_file"), mock.call("cert_key_file")]
)
def test_configuration_schema(self):
configuration_schema = Prometheus.configuration_schema()
assert configuration_schema == {
"type": "object",
"properties": {
"url": {"type": "string", "title": "Prometheus API URL"},
"verify_ssl": {
"type": "boolean",
"title": "Verify SSL (Ignored, if SSL Root Certificate is given)",
"default": True,
},
"cert_File": {"type": "string", "title": "SSL Client Certificate", "default": None},
"cert_key_File": {"type": "string", "title": "SSL Client Key", "default": None},
"ca_cert_File": {"type": "string", "title": "SSL Root Certificate", "default": None},
},
"required": ["url"],
"secret": ["cert_File", "cert_key_File", "ca_cert_File"],
"extra_options": ["verify_ssl", "cert_File", "cert_key_File", "ca_cert_File"],
}
def test_enabled(self):
assert Prometheus.enabled() is True
@mock.patch("redash.query_runner.prometheus.requests.get")
@mock.patch("redash.query_runner.prometheus.Prometheus._cleanup_cert_files")
def test_test_connection(
self,
cleanup_cert_files_mock: mock.MagicMock,
requests_get_mock: mock.MagicMock,
):
# 1. case: successful test connection
prometheus = Prometheus({"url": "url"})
prometheus_kwargs = {"verify": True, "cert": ()}
requests_get_mock.get.return_value = mock.Mock(ok=True)
connected = prometheus.test_connection()
self.assertTrue(connected)
requests_get_mock.assert_called_once_with("url", **prometheus_kwargs)
cleanup_cert_files_mock.assert_called_once_with(prometheus_kwargs)
cleanup_cert_files_mock.reset_mock()
requests_get_mock.reset_mock()
# 2. case: unsuccessful test connection
prometheus = Prometheus({"url": "url"})
prometheus_kwargs = {"verify": True, "cert": ()}
requests_get_mock.return_value = mock.Mock(ok=False)
connected = prometheus.test_connection()
self.assertFalse(connected)
requests_get_mock.assert_called_once_with("url", **prometheus_kwargs)
cleanup_cert_files_mock.assert_called_once_with(prometheus_kwargs)
cleanup_cert_files_mock.reset_mock()
requests_get_mock.reset_mock()
# 3. case: unsuccessful test connection with raised exception
prometheus = Prometheus({"url": "url"})
prometheus_kwargs = {"verify": True, "cert": ()}
requests_get_mock.side_effect = Exception("test exception")
with self.assertRaises(Exception) as exception_obj:
connected = prometheus.test_connection()
self.assertFalse(connected)
self.assertEqual(str(exception_obj.exception), "test exception")
requests_get_mock.assert_called_once_with("url", **prometheus_kwargs)
cleanup_cert_files_mock.assert_called_once_with(prometheus_kwargs)
@mock.patch("redash.query_runner.prometheus.requests.get")
@mock.patch("redash.query_runner.prometheus.Prometheus._cleanup_cert_files")
def test_get_schema(
self,
cleanup_cert_files_mock: mock.MagicMock,
requests_get_mock: mock.MagicMock,
):
# 1. case: successful get schema
prometheus = Prometheus({"url": "url"})
prometheus_kwargs = {"verify": True, "cert": ()}
requests_get_mock.return_value = mock.Mock(json=mock.Mock(return_value={"data": ["name1", "name2"]}))
schema = prometheus.get_schema()
self.assertEqual(schema, [{"name": "name1", "columns": []}, {"name": "name2", "columns": []}])
requests_get_mock.assert_called_once_with("url/api/v1/label/__name__/values", **prometheus_kwargs)
cleanup_cert_files_mock.assert_called_once_with(prometheus_kwargs)
cleanup_cert_files_mock.reset_mock()
requests_get_mock.reset_mock()
# 2. case: successful get empty schema
prometheus = Prometheus({"url": "url"})
prometheus_kwargs = {"verify": True, "cert": ()}
requests_get_mock.return_value = mock.Mock(json=mock.Mock(return_value={"data": []}))
schema = prometheus.get_schema()
self.assertEqual(schema, [])
requests_get_mock.assert_called_once_with("url/api/v1/label/__name__/values", **prometheus_kwargs)
cleanup_cert_files_mock.assert_called_once_with(prometheus_kwargs)
cleanup_cert_files_mock.reset_mock()
requests_get_mock.reset_mock()
# 3. case: unsuccessful get schema with an exception
prometheus = Prometheus({"url": "url"})
prometheus_kwargs = {"verify": True, "cert": ()}
requests_get_mock.side_effect = Exception("test exception")
with self.assertRaises(Exception) as exception_obj:
schema = prometheus.get_schema()
self.assertEqual(schema, [])
self.assertEqual(str(exception_obj.exception), "test exception")
requests_get_mock.assert_called_once_with("url/api/v1/label/__name__/values", **prometheus_kwargs)
cleanup_cert_files_mock.assert_called_once_with(prometheus_kwargs)
@mock.patch("redash.query_runner.prometheus.requests.get")
@mock.patch("redash.query_runner.prometheus.Prometheus._cleanup_cert_files")
def test_run_query(
self,
cleanup_cert_files_mock: mock.MagicMock,
requests_get_mock: mock.MagicMock,
):
# 1. case: successful run instant query
prometheus = Prometheus({"url": "url"})
prometheus_kwargs = {"verify": True, "cert": ()}
timestamp_expected_7400 = datetime.fromtimestamp(1516937400.781)
rows = [
{
"name": "example_metric_name",
"foo_bar": "foo",
"timestamp": timestamp_expected_7400,
"value": "7400_foo",
},
{
"name": "example_metric_name",
"foo_bar": "bar",
"timestamp": timestamp_expected_7400,
"value": "7400_bar",
},
]
columns = [
{"friendly_name": "timestamp", "type": "datetime", "name": "timestamp"},
{"friendly_name": "value", "type": "string", "name": "value"},
{"friendly_name": "name", "type": "string", "name": "name"},
{"friendly_name": "foo_bar", "type": "string", "name": "foo_bar"},
]
data_expected = json_dumps({"rows": rows, "columns": columns})
requests_get_mock.return_value = mock.Mock(
json=mock.Mock(return_value={"data": {"result": self.instant_query_result}})
)
data, error = prometheus.run_query("http_requests_total", "user")
self.assertEqual(data, data_expected)
self.assertIsNone(error)
requests_get_mock.assert_called_once_with(
"url/api/v1/query", params={"query": ["http_requests_total"]}, **prometheus_kwargs
)
cleanup_cert_files_mock.assert_called_once_with(prometheus_kwargs)
cleanup_cert_files_mock.reset_mock()
requests_get_mock.reset_mock()
# 2. case: successful run instant query with empty result
prometheus = Prometheus({"url": "url"})
prometheus_kwargs = {"verify": True, "cert": ()}
requests_get_mock.return_value = mock.Mock(json=mock.Mock(return_value={"data": {"result": []}}))
data, error = prometheus.run_query("http_requests_total", "user")
self.assertIsNone(data)
self.assertEqual(error, "query result is empty.")
requests_get_mock.assert_called_once_with(
"url/api/v1/query", params={"query": ["http_requests_total"]}, **prometheus_kwargs
)
cleanup_cert_files_mock.assert_called_once_with(prometheus_kwargs)
cleanup_cert_files_mock.reset_mock()
requests_get_mock.reset_mock()
# 3. case: successful run range query with start and end
prometheus = Prometheus({"url": "url"})
prometheus_kwargs = {"verify": True, "cert": ()}
timestamp_expected_8000 = datetime.fromtimestamp(1516938000.781)
rows = [
{
"name": "example_metric_name",
"foo_bar": "foo",
"timestamp": timestamp_expected_7400,
"value": "7400_foo",
},
{
"name": "example_metric_name",
"foo_bar": "foo",
"timestamp": timestamp_expected_8000,
"value": "8000_foo",
},
{
"name": "example_metric_name",
"foo_bar": "bar",
"timestamp": timestamp_expected_7400,
"value": "7400_bar",
},
{
"name": "example_metric_name",
"foo_bar": "bar",
"timestamp": timestamp_expected_8000,
"value": "8000_bar",
},
]
columns = [
{"friendly_name": "timestamp", "type": "datetime", "name": "timestamp"},
{"friendly_name": "value", "type": "string", "name": "value"},
{"friendly_name": "name", "type": "string", "name": "name"},
{"friendly_name": "foo_bar", "type": "string", "name": "foo_bar"},
]
data_expected = json_dumps({"rows": rows, "columns": columns})
requests_get_mock.return_value = mock.Mock(
json=mock.Mock(return_value={"data": {"result": self.range_query_result}})
)
start_timestamp_expected = int(time.mktime(datetime(2018, 1, 26).timetuple()))
end_timestamp_expected = int(time.mktime(datetime(2018, 1, 27).timetuple()))
data, error = prometheus.run_query(
"http_requests_total&start=2018-01-26T00:00:00.000Z&end=2018-01-27T00:00:00.000Z&step=60s", "user"
)
self.assertEqual(data, data_expected)
self.assertIsNone(error)
requests_get_mock.assert_called_once_with(
"url/api/v1/query_range",
params={
"query": ["http_requests_total"],
"start": [start_timestamp_expected],
"end": [end_timestamp_expected],
"step": ["60s"],
},
**prometheus_kwargs,
)
cleanup_cert_files_mock.assert_called_once_with(prometheus_kwargs)
cleanup_cert_files_mock.reset_mock()
requests_get_mock.reset_mock()
# 4. case: successful run range query with start and without end
prometheus = Prometheus({"url": "url"})
prometheus_kwargs = {"verify": True, "cert": ()}
rows = [
{
"name": "example_metric_name",
"foo_bar": "foo",
"timestamp": timestamp_expected_7400,
"value": "7400_foo",
},
{
"name": "example_metric_name",
"foo_bar": "foo",
"timestamp": timestamp_expected_8000,
"value": "8000_foo",
},
{
"name": "example_metric_name",
"foo_bar": "bar",
"timestamp": timestamp_expected_7400,
"value": "7400_bar",
},
{
"name": "example_metric_name",
"foo_bar": "bar",
"timestamp": timestamp_expected_8000,
"value": "8000_bar",
},
]
columns = [
{"friendly_name": "timestamp", "type": "datetime", "name": "timestamp"},
{"friendly_name": "value", "type": "string", "name": "value"},
{"friendly_name": "name", "type": "string", "name": "name"},
{"friendly_name": "foo_bar", "type": "string", "name": "foo_bar"},
]
data_expected = json_dumps({"rows": rows, "columns": columns})
now_datetime = datetime(2023, 12, 12, 11, 00, 00)
end_timestamp_expected = int(time.mktime(now_datetime.timetuple()))
requests_get_mock.return_value = mock.Mock(
json=mock.Mock(return_value={"data": {"result": self.range_query_result}})
)
with mock.patch("redash.query_runner.prometheus.Prometheus._get_datetime_now") as get_datetime_now_mock:
get_datetime_now_mock.return_value = now_datetime
data, error = prometheus.run_query("http_requests_total&start=2018-01-26T00:00:00.000Z&step=60s", "user")
self.assertEqual(data, data_expected)
self.assertIsNone(error)
requests_get_mock.assert_called_once_with(
"url/api/v1/query_range",
params={
"query": ["http_requests_total"],
"start": [start_timestamp_expected],
"end": [end_timestamp_expected],
"step": ["60s"],
},
**prometheus_kwargs,
)
cleanup_cert_files_mock.assert_called_once_with(prometheus_kwargs)
cleanup_cert_files_mock.reset_mock()
requests_get_mock.reset_mock()
data = None
error = None
# 5. case: run query with exception
prometheus = Prometheus({"url": "url"})
prometheus_kwargs = {"verify": True, "cert": ()}
requests_get_mock.side_effect = Exception("test exception")
with self.assertRaises(Exception) as exception_obj:
data, error = prometheus.run_query("http_requests_total", "user")
self.assertIsNone(data)
self.assertIsNone(error)
self.assertEqual(str(exception_obj.exception), "test exception")
requests_get_mock.assert_called_once_with(
"url/api/v1/query", params={"query": ["http_requests_total"]}, **prometheus_kwargs
)
cleanup_cert_files_mock.assert_called_once_with(prometheus_kwargs)
Reference in New Issue View Git Blame Copy Permalink