mirror of
https://github.com/langgenius/dify.git
synced 2026-05-10 05:56:31 +08:00
ae898652b2
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: wangxiaolei <fatelei@gmail.com>
74 lines
2.0 KiB
Python
74 lines
2.0 KiB
Python
import os
|
|
from collections import UserDict
|
|
|
|
import pytest
|
|
from _pytest.monkeypatch import MonkeyPatch
|
|
from upstash_vector import Index
|
|
|
|
|
|
class MockIndex:
|
|
def __init__(self, url="", token=""):
|
|
self.url = url
|
|
self.token = token
|
|
self.vectors = []
|
|
|
|
def upsert(self, vectors):
|
|
for vector in vectors:
|
|
vector.score = 0.5
|
|
self.vectors.append(vector)
|
|
return {"code": 0, "msg": "operation success", "affectedCount": len(vectors)}
|
|
|
|
def fetch(self, ids):
|
|
return [vector for vector in self.vectors if vector.id in ids]
|
|
|
|
def delete(self, ids):
|
|
self.vectors = [vector for vector in self.vectors if vector.id not in ids]
|
|
return {"code": 0, "msg": "Success"}
|
|
|
|
def query(
|
|
self,
|
|
vector: None,
|
|
top_k: int = 10,
|
|
include_vectors: bool = False,
|
|
include_metadata: bool = False,
|
|
filter: str = "",
|
|
data: str | None = None,
|
|
namespace: str = "",
|
|
include_data: bool = False,
|
|
):
|
|
mock_result = []
|
|
for vector_data in self.vectors:
|
|
mock_result.append(vector_data)
|
|
return mock_result[:top_k]
|
|
|
|
def reset(self):
|
|
self.vectors = []
|
|
|
|
def info(self):
|
|
return AttrDict({"dimension": 1024})
|
|
|
|
|
|
class AttrDict(UserDict):
|
|
def __getattr__(self, item):
|
|
return self.get(item)
|
|
|
|
|
|
MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
|
|
|
|
|
|
@pytest.fixture
|
|
def setup_upstashvector_mock(request, monkeypatch: MonkeyPatch):
|
|
if MOCK:
|
|
monkeypatch.setattr(Index, "__init__", MockIndex.__init__)
|
|
monkeypatch.setattr(Index, "upsert", MockIndex.upsert)
|
|
monkeypatch.setattr(Index, "fetch", MockIndex.fetch)
|
|
monkeypatch.setattr(Index, "delete", MockIndex.delete)
|
|
monkeypatch.setattr(Index, "query", MockIndex.query)
|
|
monkeypatch.setattr(Index, "reset", MockIndex.reset)
|
|
monkeypatch.setattr(Index, "info", MockIndex.info)
|
|
|
|
yield
|
|
|
|
if MOCK:
|
|
monkeypatch.undo()
|