diff --git a/.gitignore b/.gitignore index 0f3420a..fdf6837 100644 --- a/.gitignore +++ b/.gitignore @@ -133,3 +133,5 @@ dmypy.json # Pyre type checker .pyre/ + +.DS_Store \ No newline at end of file diff --git a/src/autometrics/prometheus_url.py b/src/autometrics/prometheus_url.py index 614add6..96e4bbd 100644 --- a/src/autometrics/prometheus_url.py +++ b/src/autometrics/prometheus_url.py @@ -4,13 +4,17 @@ class Generator: - def __init__(self, functionName, moduleName): + def __init__(self, functionName, moduleName, baseUrl=None): load_dotenv() self.functionName = functionName self.moduleName = moduleName - self.baseUrl = os.getenv("PROMETHEUS_URL") + self.baseUrl = baseUrl or os.getenv("PROMETHEUS_URL") if self.baseUrl is None: - self.baseUrl = "http://localhost:9090/" + self.baseUrl = "http://localhost:9090" + elif self.baseUrl[-1] == "/": + self.baseUrl = self.baseUrl[ + :-1 + ] # Remove the trailing slash if there is one def createURLs(self): requestRateQuery = f'sum by (function, module) (rate (function_calls_count_total{{function="{self.functionName}",module="{self.moduleName}"}}[5m]))' @@ -30,5 +34,5 @@ def createURLs(self): def createPrometheusUrl(self, query): urlEncode = urllib.parse.quote(query) - url = f"{self.baseUrl}graph?g0.expr={urlEncode}&g0.tab=0" + url = f"{self.baseUrl}/graph?g0.expr={urlEncode}&g0.tab=0" return url diff --git a/src/test_prometheus_url.py b/src/test_prometheus_url.py new file mode 100644 index 0000000..d9b7f65 --- /dev/null +++ b/src/test_prometheus_url.py @@ -0,0 +1,34 @@ +import unittest +from autometrics.prometheus_url import Generator + + +# Defaults to localhost:9090 +class TestPrometheusUrlGeneratorDefault(unittest.TestCase): + def setUp(self): + self.generator = Generator("myFunction", "myModule") + + def test_createPrometheusUrl(self): + url = self.generator.createPrometheusUrl("myQuery") + self.assertTrue( + url.startswith("http://localhost:9090/graph?g0.expr=") + ) # Make sure the base URL is correct + self.assertIn("myQuery", url) # Make sure the query is included in the URL + + +# Creates proper urls when given a custom base URL +class TestPrometheusUrlGeneratorCustomUrl(unittest.TestCase): + def setUp(self): + self.generator = Generator( + "myFunction", "myModule", baseUrl="http://localhost:9091" + ) + + def test_createPrometheusUrl(self): + url = self.generator.createPrometheusUrl("myQuery") + self.assertTrue( + url.startswith("http://localhost:9091/graph?g0.expr=") + ) # Make sure the base URL is correct + self.assertIn("myQuery", url) # Make sure the query is included in the URL + + +if __name__ == "__main__": + unittest.main()