small refactor of opensearch dest check script

scripts/opensearch-test-helpers/destination_connector/test-ingest-opensearch-output.py

@@ -20,15 +20,21 @@
 
     initial_query = {"query": {"simple_query_string": {"fields": ["text"], "query": EXPECTED_TEXT}}}
 
+    timeout = True
+    initial_embeddings = None
     for i in range(3):
         try:
             initial_result = client.search(index="ingest-test-destination", body=initial_query)
             initial_embeddings = initial_result["hits"]["hits"][0]["_source"]["embeddings"]
+            timeout = False
             break
         except:  # noqa: E722
             print("Retrying to get initial embeddings")
             time.sleep(3)
 
+    if timeout:
+        raise TimeoutError("timed out trying to get results from opensearch")
+
     query = {"size": 1, "query": {"knn": {"embeddings": {"vector": initial_embeddings, "k": 1}}}}
 
     vector_search = client.search(index="ingest-test-destination", body=query)
@@ -41,22 +47,21 @@
             "OpenSearch dest check failed:" f"Did not find {EXPECTED_TEXT} in via vector search."
         )
 
+    timeout = True
+    count = None
     for i in range(3):
         try:
-            count = int(client.count(index="ingest-test-destination")["count"])
-            assert count == N_ELEMENTS
+            count = client.count(index="ingest-test-destination")["count"]
+            timeout = False
             break
         except:  # noqa: E722
             print("Retrying to get count")
             time.sleep(3)
 
-    try:
-        count = int(client.count(index="ingest-test-destination")["count"])
-        assert count == N_ELEMENTS
-    except AssertionError:
-        sys.exit(
-            "OpenSearch dest check failed:"
-            f"got {count} items in index, expected {N_ELEMENTS} items in index."
-        )
+    if timeout:
+        raise TimeoutError("timed out trying to get count from opensearch")
+
+    assert int(count) == N_ELEMENTS, "OpenSearch dest check failed:"
+    f"got {count} items in index, expected {N_ELEMENTS} items in index."
 
     print(f"OpenSearch destination test was successful with {count} items being uploaded.")