Skip to content

Commit

Permalink
fix: pinecone delays
Browse files Browse the repository at this point in the history
  • Loading branch information
jamescalam committed Jan 5, 2025
1 parent 1141d73 commit 4930d46
Showing 1 changed file with 16 additions and 23 deletions.
39 changes: 16 additions & 23 deletions tests/unit/test_router.py
Original file line number Diff line number Diff line change
Expand Up @@ -752,37 +752,30 @@ def test_query_filter(self, routes, index_cls, encoder_cls, router_cls):
)
if index_cls is PineconeIndex:
time.sleep(PINECONE_SLEEP) # allow for index to be populated
query_result = route_layer(text="Hello", route_filter=["Route 1"]).name

# with pytest.raises(ValueError):
# # Route 8 does not exist so should raise ValueError
# route_layer(text="Hello", route_filter=["Route 8"]).name
try:
# TODO JB: currently LocalIndex raises ValueError but others don't
# they should all behave in the same way
route_layer(text="Hello", route_filter=["Route 8"]).name
except ValueError:
assert True

assert query_result in ["Route 1"]

@pytest.mark.skipif(
os.environ.get("PINECONE_API_KEY") is None, reason="Pinecone API key required"
)
def test_query_filter_pinecone(self, routes, index_cls, encoder_cls, router_cls):
if index_cls is PineconeIndex:
encoder = encoder_cls()
pineconeindex = init_index(index_cls, index_name=encoder.__class__.__name__)
route_layer = router_cls(
encoder=encoder,
routes=routes,
index=pineconeindex,
auto_sync="local",
)
time.sleep(PINECONE_SLEEP) # allow for index to be populated
count = 0
# we allow for 5 retries to allow for index to be populated
while count < RETRY_COUNT:
query_result = route_layer(text="Hello", route_filter=["Route 1"]).name

try:
route_layer(text="Hello", route_filter=["Route 8"]).name
except ValueError:
assert True

assert query_result in ["Route 1"]
assert query_result in ["Route 1"]
break
except AssertionError:
logger.warning(
f"Query result not in expected routes, waiting for retry (try {count})"
)
count += 1
time.sleep(PINECONE_SLEEP) # allow for index to be populated

@pytest.mark.skipif(
os.environ.get("PINECONE_API_KEY") is None, reason="Pinecone API key required"
Expand Down

0 comments on commit 4930d46

Please sign in to comment.