Skip to content

Commit

Permalink
Add additional functional tests
Browse files Browse the repository at this point in the history
  • Loading branch information
mpozniak95 committed Jun 27, 2024
1 parent 4cb2415 commit 16d99e0
Show file tree
Hide file tree
Showing 2 changed files with 79 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -268,7 +268,7 @@ def _traced_aggregate(func, instance, args, kwargs):
_set_span_attribute(
span,
"redis.commands.aggregate.query",
query.query_string(),
query._query,
)
response = func(*args, **kwargs)
_set_span_attribute(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,15 @@
import redis
import redis.asyncio

from redis.exceptions import ResponseError
from redis.commands.search.indexDefinition import IndexDefinition, IndexType
from redis.commands.search.aggregation import AggregateRequest
from redis.commands.search.query import Query
from redis.commands.search.field import (
TextField,
VectorField,
)

from opentelemetry import trace
from opentelemetry.instrumentation.redis import RedisInstrumentor
from opentelemetry.semconv.trace import SpanAttributes
Expand Down Expand Up @@ -614,3 +623,72 @@ def test_get(self):
self.assertEqual(
span.attributes.get(SpanAttributes.DB_STATEMENT), "GET ?"
)


class TestRedisearchInstrument(TestBase):
def setUp(self):
super().setUp()
self.redis_client = redis.Redis(port=6379)
self.redis_client.flushall()
self.embedding_dim = 256
RedisInstrumentor().instrument(tracer_provider=self.tracer_provider)
self.prepare_data()
self.create_index()

def tearDown(self):
RedisInstrumentor().uninstrument()
super().tearDown()

def prepare_data(self):
try:
self.redis_client.ft("idx:test_vss").dropindex(True)
except ResponseError:
print("No such index")
item = {"name": "test",
"value": "test_value",
"embeddings": [0.1] * 256}
pipeline = self.redis_client.pipeline()
pipeline.json().set(f"test:001", "$", item)
res = pipeline.execute()
assert False not in res

def create_index(self):
schema = (
TextField("$.name", no_stem=True, as_name="name"),
TextField("$.value", no_stem=True, as_name="value"),
VectorField("$.embeddings",
"FLAT",
{
"TYPE": "FLOAT32",
"DIM": self.embedding_dim,
"DISTANCE_METRIC": "COSINE",
},
as_name="vector",),
)
definition = IndexDefinition(prefix=["test:"], index_type=IndexType.JSON)
res = self.redis_client.ft("idx:test_vss").create_index(fields=schema, definition=definition)
assert "OK" in str(res)

def test_redis_create_index(self):
spans = self.memory_exporter.get_finished_spans()
span = next(span for span in spans if span.name == "redis.create_index")
assert "redis.create_index.definition" in span.attributes
assert "redis.create_index.fields" in span.attributes

def test_redis_aggregate(self):
query = "*"
self.redis_client.ft("idx:test_vss").aggregate(AggregateRequest(query).load())
spans = self.memory_exporter.get_finished_spans()
span = next(span for span in spans if span.name == "redis.aggregate")
assert span.attributes.get("redis.commands.aggregate.query") == query
assert "redis.commands.aggregate.results" in span.attributes

def test_redis_query(self):
query = "@name:test"
res = self.redis_client.ft("idx:test_vss").search(Query(query))

spans = self.memory_exporter.get_finished_spans()
span = next(span for span in spans if span.name == "redis.search")

assert span.attributes.get("redis.commands.search.query") == query
assert span.attributes.get("redis.commands.search.total") == 1

0 comments on commit 16d99e0

Please sign in to comment.