Skip to content
This repository has been archived by the owner on Mar 1, 2024. It is now read-only.

Commit

Permalink
added unit test and additional metadata parameter
Browse files Browse the repository at this point in the history
  • Loading branch information
@mjimenezp
mjimenezp committed Feb 3, 2024
1 parent 8edd92d commit a7ac116
Show file tree
Hide file tree
Showing 2 changed files with 223 additions and 9 deletions.
19 changes: 10 additions & 9 deletions llama_hub/prometheus/base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,14 @@ def __init__(self, endpoint: str, size: Optional[int] = 100):


def load_data(self,
query: str,
start_time: Optional[datetime],
end_time: Optional[datetime],
step: Optional[str],
get_pararameters: Optional[dict],
metadata_fields: Optional[list]=[]
) -> List[Document]:
query: str,
start_time: Optional[datetime],
end_time: Optional[datetime],
step: Optional[str],
get_pararameters: Optional[dict],
metadata_fields: Optional[list]=[],
additional_metadata: Optional[dict]={}
) -> List[Document]:

if start_time is not None and end_time is not None:
params = {
Expand Down Expand Up @@ -53,11 +54,11 @@ def load_data(self,
if key_value in metadata_fields:
metadata[key_value] = row['metric'][key_value]
for value in row['values']:
current_metadata = metadata.copy()
current_metadata = metadata | additional_metadata
current_metadata['timestamp'] = value[0]
documents.append(
Document(text=value[1], extra_info=current_metadata, embedding=None)
)
return documents
return documents


213 changes: 213 additions & 0 deletions tests/test_prometheus/test_prometheus.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,213 @@
import pytest
from typing import Any
from pytest_mock import MockerFixture
from llama_hub.prometheus import PrometheusReader
from datetime import datetime, timedelta



def test_prometheus_reader_with_query_range(mocker) -> None:
end_time = datetime.utcnow() - timedelta(days=1)
start_time = end_time - timedelta(hours=1)
mocked_metric_data = [{
"metric":{
"endpoint": "http-metrics",
"metric_name": "test_metric",
"pod": "test-pod",
},
"values":[]
}]
for metric_info in mocked_metric_data:
for i in range(0,60,5):
timestamp = int((start_time + timedelta(minutes=i)).timestamp())
metric_info["values"].append([timestamp, i*100])

mocker.patch("prometheus_api_client.PrometheusConnect.custom_query_range",return_value = mocked_metric_data)
reader = PrometheusReader(endpoint="test-endpoint", size=100)

query = 'avg_over_time(test_metric[5m])'
metadata_fields = []


documents = reader.load_data(query=query,
start_time=start_time,
end_time=end_time,
step=None,
get_pararameters=None,
metadata_fields=metadata_fields)
assert len(documents) == 12
counter = 0
for doc in range(0,len(documents)):
assert documents[doc].text == str(doc * 5 * 100)


def test_prometheus_reader_with_query(mocker) -> None:
mocked_metric_data = [{
"metric":{
"Name": "/dev_cicd_test1",
"metric_name": "docker_memory_usage",
"host": "dl1.dockerserver.com",
"offering": "docker",
},
"values":[
[
1706898840,
"338149320.8186831"
],
[
1706898900,
"338146802.62924147"
],
[
1706898960,
"281811971.070067"
],
[
1706899020,
"281813920.81963813"
],
[
1706899080,
"225405325.11296228"
],
[
1706899140,
"281780621.1695041"
],
[
1706899200,
"281800467.21272033"
],
[
1706899260,
"338153045.11319935"
],
[
1706899320,
"338145581.4473358"
],
[
1706899380,
"338119211.1141539"
]
]
},
{
"metric":{
"Name": "/prod_cicd_test2",
"metric_name": "docker_memory_usage",
"host": "dl1.dockerserver.com",
"offering": "docker",
},
"values":[
[
1706898840,
"106213135.11232835"
],
[
1706898900,
"106212621.7800675"
],
[
1706898960,
"106221285.85705297"
],
[
1706899020,
"79666690.5925202"
],
[
1706899080,
"79650036.00729768"
],
[
1706899140,
"106208946.49888971"
],
[
1706899200,
"106216536.73648405"
],
[
1706899260,
"79660663.82172346"
],
[
1706899320,
"106211849.69262712"
],
[
1706899380,
"53101858.29009599"
]
]
}]
mocker.patch("prometheus_api_client.PrometheusConnect.custom_query",return_value = mocked_metric_data)
reader = PrometheusReader(endpoint="test-endpoint", size=100)
query = "rate(docker_memory_usage[5m])"
metadata_fields = ["Name", "metric_name", "host", "offering"]
documents = reader.load_data(query=query,
start_time=None,
end_time=None,
step=None,
get_pararameters=None,
metadata_fields=metadata_fields,
additional_metadata={})
assert len(documents) == 20
for doc in documents:
field_count = 0
for field in doc.metadata.keys():
if field in metadata_fields:
field_count += 1
assert field_count == len(metadata_fields)


def test_prometheus_reader_with_empty_result(mocker) -> None:
mocked_metric_data = []
mocker.patch("prometheus_api_client.PrometheusConnect.custom_query",return_value = mocked_metric_data)
reader = PrometheusReader(endpoint="test-endpoint", size=100)
query = "rate(non_existing_metric[1m])"
documents = reader.load_data(query=query,
start_time=None,
end_time=None,
step=None,
get_pararameters=None,
metadata_fields=[],
additional_metadata={})
assert len(documents) == 0

def test_prometheus_reader_with_additional_metadata(mocker) -> None:
end_time = datetime.utcnow() - timedelta(days=1)
start_time = end_time - timedelta(hours=1)
mocked_metric_data = [{
"metric":{
"endpoint": "http-metrics",
"metric_name": "test_metric",
"pod": "test-pod",
},
"values":[]
}]
for metric_info in mocked_metric_data:
for i in range(0,60,5):
timestamp = int((start_time + timedelta(minutes=i)).timestamp())
metric_info["values"].append([timestamp, i*100])

mocker.patch("prometheus_api_client.PrometheusConnect.custom_query_range",return_value = mocked_metric_data)
reader = PrometheusReader(endpoint="test-endpoint", size=100)
query = 'avg_over_time(test_metric[5m])'
metadata_fields = ['endpoint', 'metric_name', 'pod']
additional_metadata = {
"collector": "prometheus_reader",
"collection_date": datetime.now().isoformat(),
}
documents = reader.load_data(query=query,
start_time=start_time,
end_time=end_time,
step=None,
get_pararameters=None,
metadata_fields=metadata_fields,
additional_metadata=additional_metadata)
assert len(documents) == 12
for doc in documents:
all_items_exist = all(key in doc.metadata and doc.metadata[key] == value for key, value in additional_metadata.items())
assert all_items_exist

0 comments on commit a7ac116

Please sign in to comment.