-
Notifications
You must be signed in to change notification settings - Fork 4
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add integration test for vectorstore
Signed-off-by: B-Step62 <[email protected]>
- Loading branch information
Showing
3 changed files
with
96 additions
and
0 deletions.
There are no files selected for viewing
11 changes: 11 additions & 0 deletions
11
libs/databricks/tests/integration_tests/test_chat_models.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
12 changes: 12 additions & 0 deletions
12
libs/databricks/tests/integration_tests/test_embeddings.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
73 changes: 73 additions & 0 deletions
73
libs/databricks/tests/integration_tests/test_vectorstore.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,73 @@ | ||
""" | ||
This file contains the integration test for DatabricksVectorSearch class. | ||
We run the integration tests nightly by the trusted CI/CD system defined in | ||
a private repository, in order to securely run the tests. With this design, | ||
integration test is not intended to be run manually by OSS contributors. | ||
If you want to update the DatabricksVectorSearch implementation and you think | ||
that you need to update the corresponding integration test, please contact to | ||
the maintainers of the repository to verify the changes. | ||
""" | ||
|
||
import os | ||
import time | ||
|
||
import requests | ||
|
||
|
||
def test_vectorstore(): | ||
""" | ||
We run the integration tests for vector store by Databricks Workflow, | ||
because the setup is too complex to run within a single python file. | ||
Thereby, this test simply triggers the workflow by calling the REST API. | ||
""" | ||
test_endpoint = os.getenv("DATABRICKS_HOST") | ||
test_job_id = os.getenv("VS_TEST_JOB_ID") | ||
headers = { | ||
"Authorization": f"Bearer {os.getenv('DATABRICKS_TOKEN')}", | ||
} | ||
|
||
# Check if there is any ongoing job run | ||
response = requests.get( | ||
f"{test_endpoint}/api/2.1/jobs/runs/list", | ||
json={ | ||
"job_id": test_job_id, | ||
"active_only": True, | ||
}, | ||
headers=headers, | ||
) | ||
no_active_run = "runs" not in response.json() or len(response.json()["runs"]) == 0 | ||
assert no_active_run, "There is an ongoing job run. Please wait for it to complete." | ||
|
||
# Trigger the workflow | ||
response = requests.post( | ||
f"{test_endpoint}/api/2.1/jobs/run-now", | ||
json={ | ||
"job_id": test_job_id, | ||
}, | ||
headers=headers, | ||
) | ||
|
||
assert response.status_code == 200, "Failed to trigger the workflow." | ||
|
||
# Wait for the job to complete | ||
while True: | ||
response = requests.get( | ||
f"{test_endpoint}/api/2.1/jobs/runs/get", | ||
json={ | ||
"run_id": response.json()["run_id"], | ||
}, | ||
headers=headers, | ||
) | ||
|
||
assert response.status_code == 200, "Failed to get the job status." | ||
|
||
status = response.json()["status"] | ||
if status["state"] == "TERMINATED": | ||
if status["termination_details"]["type"] == "SUCCESS": | ||
break | ||
else: | ||
assert False, "Job failed. Please check the logs in the workspace." | ||
|
||
time.sleep(60) | ||
print("Job is still running...") # noqa: T201 |