admin/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

chore(api/tests): apply ruff reformat #7590 (#7591)

Browse Source 
Co-authored-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
Bowen Liang
2024-08-23 23:52:25 +08:00
committed by GitHub
parent 2da63654e5
commit b035c02f78
155 changed files with 4279 additions and 5925 deletions
Download Patch File Download Diff File Expand all files Collapse all files

109
api/tests/integration_tests/vdb/opensearch/test_opensearch.py
View File

@@ -29,54 +29,55 @@ class TestOpenSearchVector:
self.example_doc_id = "example_doc_id"
self.vector = OpenSearchVector(
collection_name=self.collection_name,
config=OpenSearchConfig(
host='localhost',
port=9200,
user='admin',
password='password',
secure=False
)
config=OpenSearchConfig(host="localhost", port=9200, user="admin", password="password", secure=False),
)
self.vector._client = MagicMock()
@pytest.mark.parametrize("search_response, expected_length, expected_doc_id", [
({
'hits': {
'total': {'value': 1},
'hits': [
{'_source': {'page_content': get_example_text(), 'metadata': {"document_id": "example_doc_id"}}}
]
}
}, 1, "example_doc_id"),
({
'hits': {
'total': {'value': 0},
'hits': []
}
}, 0, None)
])
@pytest.mark.parametrize(
"search_response, expected_length, expected_doc_id",
[
(
{
"hits": {
"total": {"value": 1},
"hits": [
{
"_source": {
"page_content": get_example_text(),
"metadata": {"document_id": "example_doc_id"},
}
}
],
}
},
1,
"example_doc_id",
),
({"hits": {"total": {"value": 0}, "hits": []}}, 0, None),
],
)
def test_search_by_full_text(self, search_response, expected_length, expected_doc_id):
self.vector._client.search.return_value = search_response
hits_by_full_text = self.vector.search_by_full_text(query=get_example_text())
assert len(hits_by_full_text) == expected_length
if expected_length > 0:
assert hits_by_full_text[0].metadata['document_id'] == expected_doc_id
assert hits_by_full_text[0].metadata["document_id"] == expected_doc_id
def test_search_by_vector(self):
vector = [0.1] * 128
mock_response = {
'hits': {
'total': {'value': 1},
'hits': [
"hits": {
"total": {"value": 1},
"hits": [
{
'_source': {
"_source": {
Field.CONTENT_KEY.value: get_example_text(),
Field.METADATA_KEY.value: {"document_id": self.example_doc_id}
Field.METADATA_KEY.value: {"document_id": self.example_doc_id},
},
'_score': 1.0
"_score": 1.0,
}
]
],
}
}
self.vector._client.search.return_value = mock_response
@@ -85,53 +86,45 @@ class TestOpenSearchVector:
print("Hits by vector:", hits_by_vector)
print("Expected document ID:", self.example_doc_id)
print("Actual document ID:", hits_by_vector[0].metadata['document_id'] if hits_by_vector else "No hits")
print("Actual document ID:", hits_by_vector[0].metadata["document_id"] if hits_by_vector else "No hits")
assert len(hits_by_vector) > 0, f"Expected at least one hit, got {len(hits_by_vector)}"
assert hits_by_vector[0].metadata['document_id'] == self.example_doc_id, \
f"Expected document ID {self.example_doc_id}, got {hits_by_vector[0].metadata['document_id']}"
assert (
hits_by_vector[0].metadata["document_id"] == self.example_doc_id
), f"Expected document ID {self.example_doc_id}, got {hits_by_vector[0].metadata['document_id']}"
def test_get_ids_by_metadata_field(self):
mock_response = {
'hits': {
'total': {'value': 1},
'hits': [{'_id': 'mock_id'}]
}
}
mock_response = {"hits": {"total": {"value": 1}, "hits": [{"_id": "mock_id"}]}}
self.vector._client.search.return_value = mock_response
doc = Document(page_content="Test content", metadata={"document_id": self.example_doc_id})
embedding = [0.1] * 128
with patch('opensearchpy.helpers.bulk') as mock_bulk:
with patch("opensearchpy.helpers.bulk") as mock_bulk:
mock_bulk.return_value = ([], [])
self.vector.add_texts([doc], [embedding])
ids = self.vector.get_ids_by_metadata_field(key='document_id', value=self.example_doc_id)
ids = self.vector.get_ids_by_metadata_field(key="document_id", value=self.example_doc_id)
assert len(ids) == 1
assert ids[0] == 'mock_id'
assert ids[0] == "mock_id"
def test_add_texts(self):
self.vector._client.index.return_value = {'result': 'created'}
self.vector._client.index.return_value = {"result": "created"}
doc = Document(page_content="Test content", metadata={"document_id": self.example_doc_id})
embedding = [0.1] * 128
with patch('opensearchpy.helpers.bulk') as mock_bulk:
with patch("opensearchpy.helpers.bulk") as mock_bulk:
mock_bulk.return_value = ([], [])
self.vector.add_texts([doc], [embedding])
mock_response = {
'hits': {
'total': {'value': 1},
'hits': [{'_id': 'mock_id'}]
}
}
mock_response = {"hits": {"total": {"value": 1}, "hits": [{"_id": "mock_id"}]}}
self.vector._client.search.return_value = mock_response
ids = self.vector.get_ids_by_metadata_field(key='document_id', value=self.example_doc_id)
ids = self.vector.get_ids_by_metadata_field(key="document_id", value=self.example_doc_id)
assert len(ids) == 1
assert ids[0] == 'mock_id'
assert ids[0] == "mock_id"
@pytest.mark.usefixtures("setup_mock_redis")
class TestOpenSearchVectorWithRedis:
@@ -141,11 +134,11 @@ class TestOpenSearchVectorWithRedis:
def test_search_by_full_text(self):
self.tester.setup_method()
search_response = {
'hits': {
'total': {'value': 1},
'hits': [
{'_source': {'page_content': get_example_text(), 'metadata': {"document_id": "example_doc_id"}}}
]
"hits": {
"total": {"value": 1},
"hits": [
{"_source": {"page_content": get_example_text(), "metadata": {"document_id": "example_doc_id"}}}
],
}
}
expected_length = 1