Merge branch 'main' into fix/chore-fix

This commit is contained in:
Yeuoly
2024-11-01 16:23:04 +08:00
396 changed files with 15138 additions and 2610 deletions

View File

@ -91,3 +91,10 @@ INNER_API_KEY=
# Marketplace configuration
MARKETPLACE_API_URL=
# VESSL AI Credentials
VESSL_AI_MODEL_NAME=
VESSL_AI_API_KEY=
VESSL_AI_ENDPOINT_URL=
# Gitee AI Credentials
GITEE_AI_API_KEY=

View File

@ -0,0 +1,132 @@
import os
from collections.abc import Generator
import pytest
from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta
from core.model_runtime.entities.message_entities import (
AssistantPromptMessage,
PromptMessageTool,
SystemPromptMessage,
UserPromptMessage,
)
from core.model_runtime.entities.model_entities import AIModelEntity
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.gitee_ai.llm.llm import GiteeAILargeLanguageModel
def test_predefined_models():
model = GiteeAILargeLanguageModel()
model_schemas = model.predefined_models()
assert len(model_schemas) >= 1
assert isinstance(model_schemas[0], AIModelEntity)
def test_validate_credentials_for_chat_model():
model = GiteeAILargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError):
# model name to gpt-3.5-turbo because of mocking
model.validate_credentials(model="gpt-3.5-turbo", credentials={"api_key": "invalid_key"})
model.validate_credentials(
model="Qwen2-7B-Instruct",
credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")},
)
def test_invoke_chat_model():
model = GiteeAILargeLanguageModel()
result = model.invoke(
model="Qwen2-7B-Instruct",
credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Hello World!"),
],
model_parameters={
"temperature": 0.0,
"top_p": 1.0,
"presence_penalty": 0.0,
"frequency_penalty": 0.0,
"max_tokens": 10,
"stream": False,
},
stop=["How"],
stream=False,
user="foo",
)
assert isinstance(result, LLMResult)
assert len(result.message.content) > 0
def test_invoke_stream_chat_model():
model = GiteeAILargeLanguageModel()
result = model.invoke(
model="Qwen2-7B-Instruct",
credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Hello World!"),
],
model_parameters={"temperature": 0.0, "max_tokens": 100, "stream": False},
stream=True,
user="foo",
)
assert isinstance(result, Generator)
for chunk in result:
assert isinstance(chunk, LLMResultChunk)
assert isinstance(chunk.delta, LLMResultChunkDelta)
assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
if chunk.delta.finish_reason is not None:
assert chunk.delta.usage is not None
def test_get_num_tokens():
model = GiteeAILargeLanguageModel()
num_tokens = model.get_num_tokens(
model="Qwen2-7B-Instruct",
credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")},
prompt_messages=[UserPromptMessage(content="Hello World!")],
)
assert num_tokens == 10
num_tokens = model.get_num_tokens(
model="Qwen2-7B-Instruct",
credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Hello World!"),
],
tools=[
PromptMessageTool(
name="get_weather",
description="Determine weather in my location",
parameters={
"type": "object",
"properties": {
"location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},
"unit": {"type": "string", "enum": ["c", "f"]},
},
"required": ["location"],
},
),
],
)
assert num_tokens == 77

View File

@ -0,0 +1,15 @@
import os
import pytest
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.gitee_ai.gitee_ai import GiteeAIProvider
def test_validate_provider_credentials():
provider = GiteeAIProvider()
with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials(credentials={"api_key": "invalid_key"})
provider.validate_provider_credentials(credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")})

View File

@ -0,0 +1,47 @@
import os
import pytest
from core.model_runtime.entities.rerank_entities import RerankResult
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.gitee_ai.rerank.rerank import GiteeAIRerankModel
def test_validate_credentials():
model = GiteeAIRerankModel()
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(
model="bge-reranker-v2-m3",
credentials={"api_key": "invalid_key"},
)
model.validate_credentials(
model="bge-reranker-v2-m3",
credentials={
"api_key": os.environ.get("GITEE_AI_API_KEY"),
},
)
def test_invoke_model():
model = GiteeAIRerankModel()
result = model.invoke(
model="bge-reranker-v2-m3",
credentials={
"api_key": os.environ.get("GITEE_AI_API_KEY"),
},
query="What is the capital of the United States?",
docs=[
"Carson City is the capital city of the American state of Nevada. At the 2010 United States "
"Census, Carson City had a population of 55,274.",
"The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that "
"are a political division controlled by the United States. Its capital is Saipan.",
],
top_n=1,
score_threshold=0.01,
)
assert isinstance(result, RerankResult)
assert len(result.docs) == 1
assert result.docs[0].score >= 0.01

View File

@ -0,0 +1,45 @@
import os
import pytest
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.gitee_ai.speech2text.speech2text import GiteeAISpeech2TextModel
def test_validate_credentials():
model = GiteeAISpeech2TextModel()
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(
model="whisper-base",
credentials={"api_key": "invalid_key"},
)
model.validate_credentials(
model="whisper-base",
credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")},
)
def test_invoke_model():
model = GiteeAISpeech2TextModel()
# Get the directory of the current file
current_dir = os.path.dirname(os.path.abspath(__file__))
# Get assets directory
assets_dir = os.path.join(os.path.dirname(current_dir), "assets")
# Construct the path to the audio file
audio_file_path = os.path.join(assets_dir, "audio.mp3")
# Open the file and get the file object
with open(audio_file_path, "rb") as audio_file:
file = audio_file
result = model.invoke(
model="whisper-base", credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")}, file=file
)
assert isinstance(result, str)
assert result == "1 2 3 4 5 6 7 8 9 10"

View File

@ -0,0 +1,46 @@
import os
import pytest
from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.gitee_ai.text_embedding.text_embedding import GiteeAIEmbeddingModel
def test_validate_credentials():
model = GiteeAIEmbeddingModel()
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(model="bge-large-zh-v1.5", credentials={"api_key": "invalid_key"})
model.validate_credentials(model="bge-large-zh-v1.5", credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")})
def test_invoke_model():
model = GiteeAIEmbeddingModel()
result = model.invoke(
model="bge-large-zh-v1.5",
credentials={
"api_key": os.environ.get("GITEE_AI_API_KEY"),
},
texts=["hello", "world"],
user="user",
)
assert isinstance(result, TextEmbeddingResult)
assert len(result.embeddings) == 2
def test_get_num_tokens():
model = GiteeAIEmbeddingModel()
num_tokens = model.get_num_tokens(
model="bge-large-zh-v1.5",
credentials={
"api_key": os.environ.get("GITEE_AI_API_KEY"),
},
texts=["hello", "world"],
)
assert num_tokens == 2

View File

@ -0,0 +1,23 @@
import os
from core.model_runtime.model_providers.gitee_ai.tts.tts import GiteeAIText2SpeechModel
def test_invoke_model():
model = GiteeAIText2SpeechModel()
result = model.invoke(
model="speecht5_tts",
tenant_id="test",
credentials={
"api_key": os.environ.get("GITEE_AI_API_KEY"),
},
content_text="Hello, world!",
voice="",
)
content = b""
for chunk in result:
content += chunk
assert content != b""

View File

@ -0,0 +1,131 @@
import os
from collections.abc import Generator
import pytest
from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta
from core.model_runtime.entities.message_entities import (
AssistantPromptMessage,
SystemPromptMessage,
UserPromptMessage,
)
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.vessl_ai.llm.llm import VesslAILargeLanguageModel
def test_validate_credentials():
model = VesslAILargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(
model=os.environ.get("VESSL_AI_MODEL_NAME"),
credentials={
"api_key": "invalid_key",
"endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
"mode": "chat",
},
)
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(
model=os.environ.get("VESSL_AI_MODEL_NAME"),
credentials={
"api_key": os.environ.get("VESSL_AI_API_KEY"),
"endpoint_url": "http://invalid_url",
"mode": "chat",
},
)
model.validate_credentials(
model=os.environ.get("VESSL_AI_MODEL_NAME"),
credentials={
"api_key": os.environ.get("VESSL_AI_API_KEY"),
"endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
"mode": "chat",
},
)
def test_invoke_model():
model = VesslAILargeLanguageModel()
response = model.invoke(
model=os.environ.get("VESSL_AI_MODEL_NAME"),
credentials={
"api_key": os.environ.get("VESSL_AI_API_KEY"),
"endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
"mode": "chat",
},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Who are you?"),
],
model_parameters={
"temperature": 1.0,
"top_k": 2,
"top_p": 0.5,
},
stop=["How"],
stream=False,
user="abc-123",
)
assert isinstance(response, LLMResult)
assert len(response.message.content) > 0
def test_invoke_stream_model():
model = VesslAILargeLanguageModel()
response = model.invoke(
model=os.environ.get("VESSL_AI_MODEL_NAME"),
credentials={
"api_key": os.environ.get("VESSL_AI_API_KEY"),
"endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
"mode": "chat",
},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Who are you?"),
],
model_parameters={
"temperature": 1.0,
"top_k": 2,
"top_p": 0.5,
},
stop=["How"],
stream=True,
user="abc-123",
)
assert isinstance(response, Generator)
for chunk in response:
assert isinstance(chunk, LLMResultChunk)
assert isinstance(chunk.delta, LLMResultChunkDelta)
assert isinstance(chunk.delta.message, AssistantPromptMessage)
def test_get_num_tokens():
model = VesslAILargeLanguageModel()
num_tokens = model.get_num_tokens(
model=os.environ.get("VESSL_AI_MODEL_NAME"),
credentials={
"api_key": os.environ.get("VESSL_AI_API_KEY"),
"endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Hello World!"),
],
)
assert isinstance(num_tokens, int)
assert num_tokens == 21

View File

@ -0,0 +1,75 @@
import os
from typing import Optional
import pytest
from _pytest.monkeypatch import MonkeyPatch
from upstash_vector import Index
# Mocking the Index class from upstash_vector
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: Optional[str] = None,
namespace: str = "",
include_data: bool = False,
):
# Simple mock query, in real scenario you would calculate similarity
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(dict):
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()

View File

@ -0,0 +1,50 @@
import subprocess
import time
from core.rag.datasource.vdb.couchbase.couchbase_vector import CouchbaseConfig, CouchbaseVector
from tests.integration_tests.vdb.test_vector_store import (
AbstractVectorTest,
get_example_text,
setup_mock_redis,
)
def wait_for_healthy_container(service_name="couchbase-server", timeout=300):
start_time = time.time()
while time.time() - start_time < timeout:
result = subprocess.run(
["docker", "inspect", "--format", "{{.State.Health.Status}}", service_name], capture_output=True, text=True
)
if result.stdout.strip() == "healthy":
print(f"{service_name} is healthy!")
return True
else:
print(f"Waiting for {service_name} to be healthy...")
time.sleep(10)
raise TimeoutError(f"{service_name} did not become healthy in time")
class CouchbaseTest(AbstractVectorTest):
def __init__(self):
super().__init__()
self.vector = CouchbaseVector(
collection_name=self.collection_name,
config=CouchbaseConfig(
connection_string="couchbase://127.0.0.1",
user="Administrator",
password="password",
bucket_name="Embeddings",
scope_name="_default",
),
)
def search_by_vector(self):
# brief sleep to ensure document is indexed
time.sleep(5)
hits_by_vector = self.vector.search_by_vector(query_vector=self.example_embedding)
assert len(hits_by_vector) == 1
def test_couchbase(setup_mock_redis):
wait_for_healthy_container("couchbase-server", timeout=60)
CouchbaseTest().run_all_tests()

View File

@ -0,0 +1,71 @@
from unittest.mock import MagicMock, patch
import pytest
from core.rag.datasource.vdb.oceanbase.oceanbase_vector import (
OceanBaseVector,
OceanBaseVectorConfig,
)
from tests.integration_tests.vdb.__mock.tcvectordb import setup_tcvectordb_mock
from tests.integration_tests.vdb.test_vector_store import (
AbstractVectorTest,
get_example_text,
setup_mock_redis,
)
@pytest.fixture
def oceanbase_vector():
return OceanBaseVector(
"dify_test_collection",
config=OceanBaseVectorConfig(
host="127.0.0.1",
port="2881",
user="root@test",
database="test",
password="test",
),
)
class OceanBaseVectorTest(AbstractVectorTest):
def __init__(self, vector: OceanBaseVector):
super().__init__()
self.vector = vector
def search_by_vector(self):
hits_by_vector = self.vector.search_by_vector(query_vector=self.example_embedding)
assert len(hits_by_vector) == 0
def search_by_full_text(self):
hits_by_full_text = self.vector.search_by_full_text(query=get_example_text())
assert len(hits_by_full_text) == 0
def text_exists(self):
exist = self.vector.text_exists(self.example_doc_id)
assert exist == True
def get_ids_by_metadata_field(self):
ids = self.vector.get_ids_by_metadata_field(key="document_id", value=self.example_doc_id)
assert len(ids) == 0
@pytest.fixture
def setup_mock_oceanbase_client():
with patch("core.rag.datasource.vdb.oceanbase.oceanbase_vector.ObVecClient", new_callable=MagicMock) as mock_client:
yield mock_client
@pytest.fixture
def setup_mock_oceanbase_vector(oceanbase_vector):
with patch.object(oceanbase_vector, "_client"):
yield oceanbase_vector
def test_oceanbase_vector(
setup_mock_redis,
setup_mock_oceanbase_client,
setup_mock_oceanbase_vector,
oceanbase_vector,
):
OceanBaseVectorTest(oceanbase_vector).run_all_tests()

View File

@ -0,0 +1,28 @@
from core.rag.datasource.vdb.upstash.upstash_vector import UpstashVector, UpstashVectorConfig
from core.rag.models.document import Document
from tests.integration_tests.vdb.__mock.upstashvectordb import setup_upstashvector_mock
from tests.integration_tests.vdb.test_vector_store import AbstractVectorTest, get_example_text
class UpstashVectorTest(AbstractVectorTest):
def __init__(self):
super().__init__()
self.vector = UpstashVector(
collection_name="test_collection",
config=UpstashVectorConfig(
url="your-server-url",
token="your-access-token",
),
)
def get_ids_by_metadata_field(self):
ids = self.vector.get_ids_by_metadata_field(key="document_id", value=self.example_doc_id)
assert len(ids) != 0
def search_by_full_text(self):
hits_by_full_text: list[Document] = self.vector.search_by_full_text(query=get_example_text())
assert len(hits_by_full_text) == 0
def test_upstash_vector(setup_upstashvector_mock):
UpstashVectorTest().run_all_tests()

View File

@ -430,3 +430,37 @@ def test_multi_colons_parse(setup_http_mock):
assert urlencode({"Redirect": "http://example2.com"}) in result.process_data.get("request", "")
assert 'form-data; name="Redirect"\r\n\r\nhttp://example6.com' in result.process_data.get("request", "")
# assert "http://example3.com" == resp.get("headers", {}).get("referer")
def test_image_file(monkeypatch):
from types import SimpleNamespace
monkeypatch.setattr(
"core.tools.tool_file_manager.ToolFileManager.create_file_by_raw",
lambda *args, **kwargs: SimpleNamespace(id="1"),
)
node = init_http_node(
config={
"id": "1",
"data": {
"title": "http",
"desc": "",
"method": "get",
"url": "https://cloud.dify.ai/logo/logo-site.png",
"authorization": {
"type": "no-auth",
"config": None,
},
"params": "",
"headers": "",
"body": None,
},
}
)
result = node._run()
assert result.process_data is not None
assert result.outputs is not None
resp = result.outputs
assert len(resp.get("files", [])) == 1

View File

@ -22,17 +22,3 @@ from controllers.console.version import _has_new_version
)
def test_has_new_version(latest_version, current_version, expected):
assert _has_new_version(latest_version=latest_version, current_version=current_version) == expected
def test_has_new_version_invalid_input():
with pytest.raises(ValueError):
_has_new_version(latest_version="1.0", current_version="1.0.0")
with pytest.raises(ValueError):
_has_new_version(latest_version="1.0.0", current_version="1.0")
with pytest.raises(ValueError):
_has_new_version(latest_version="invalid", current_version="1.0.0")
with pytest.raises(ValueError):
_has_new_version(latest_version="1.0.0", current_version="invalid")

View File

@ -0,0 +1,125 @@
import pytest
from core.app.entities.app_invoke_entities import InvokeFrom
from core.file import File, FileTransferMethod, FileType
from core.model_runtime.entities.message_entities import ImagePromptMessageContent
from core.variables import ArrayAnySegment, ArrayFileSegment, NoneSegment
from core.workflow.entities.variable_pool import VariablePool
from core.workflow.graph_engine import Graph, GraphInitParams, GraphRuntimeState
from core.workflow.nodes.answer import AnswerStreamGenerateRoute
from core.workflow.nodes.end import EndStreamParam
from core.workflow.nodes.llm.entities import ContextConfig, LLMNodeData, ModelConfig, VisionConfig, VisionConfigOptions
from core.workflow.nodes.llm.node import LLMNode
from models.enums import UserFrom
from models.workflow import WorkflowType
class TestLLMNode:
@pytest.fixture
def llm_node(self):
data = LLMNodeData(
title="Test LLM",
model=ModelConfig(provider="openai", name="gpt-3.5-turbo", mode="chat", completion_params={}),
prompt_template=[],
memory=None,
context=ContextConfig(enabled=False),
vision=VisionConfig(
enabled=True,
configs=VisionConfigOptions(
variable_selector=["sys", "files"],
detail=ImagePromptMessageContent.DETAIL.HIGH,
),
),
)
variable_pool = VariablePool(
system_variables={},
user_inputs={},
)
node = LLMNode(
id="1",
config={
"id": "1",
"data": data.model_dump(),
},
graph_init_params=GraphInitParams(
tenant_id="1",
app_id="1",
workflow_type=WorkflowType.WORKFLOW,
workflow_id="1",
graph_config={},
user_id="1",
user_from=UserFrom.ACCOUNT,
invoke_from=InvokeFrom.SERVICE_API,
call_depth=0,
),
graph=Graph(
root_node_id="1",
answer_stream_generate_routes=AnswerStreamGenerateRoute(
answer_dependencies={},
answer_generate_route={},
),
end_stream_param=EndStreamParam(
end_dependencies={},
end_stream_variable_selector_mapping={},
),
),
graph_runtime_state=GraphRuntimeState(
variable_pool=variable_pool,
start_at=0,
),
)
return node
def test_fetch_files_with_file_segment(self, llm_node):
file = File(
id="1",
tenant_id="test",
type=FileType.IMAGE,
filename="test.jpg",
transfer_method=FileTransferMethod.LOCAL_FILE,
related_id="1",
)
llm_node.graph_runtime_state.variable_pool.add(["sys", "files"], file)
result = llm_node._fetch_files(selector=["sys", "files"])
assert result == [file]
def test_fetch_files_with_array_file_segment(self, llm_node):
files = [
File(
id="1",
tenant_id="test",
type=FileType.IMAGE,
filename="test1.jpg",
transfer_method=FileTransferMethod.LOCAL_FILE,
related_id="1",
),
File(
id="2",
tenant_id="test",
type=FileType.IMAGE,
filename="test2.jpg",
transfer_method=FileTransferMethod.LOCAL_FILE,
related_id="2",
),
]
llm_node.graph_runtime_state.variable_pool.add(["sys", "files"], ArrayFileSegment(value=files))
result = llm_node._fetch_files(selector=["sys", "files"])
assert result == files
def test_fetch_files_with_none_segment(self, llm_node):
llm_node.graph_runtime_state.variable_pool.add(["sys", "files"], NoneSegment())
result = llm_node._fetch_files(selector=["sys", "files"])
assert result == []
def test_fetch_files_with_array_any_segment(self, llm_node):
llm_node.graph_runtime_state.variable_pool.add(["sys", "files"], ArrayAnySegment(value=[]))
result = llm_node._fetch_files(selector=["sys", "files"])
assert result == []
def test_fetch_files_with_non_existent_variable(self, llm_node):
result = llm_node._fetch_files(selector=["sys", "files"])
assert result == []

View File

@ -63,17 +63,24 @@ def test_run_invalid_variable_type(document_extractor_node, mock_graph_runtime_s
@pytest.mark.parametrize(
("mime_type", "file_content", "expected_text", "transfer_method"),
("mime_type", "file_content", "expected_text", "transfer_method", "extension"),
[
("text/plain", b"Hello, world!", ["Hello, world!"], FileTransferMethod.LOCAL_FILE),
("application/pdf", b"%PDF-1.5\n%Test PDF content", ["Mocked PDF content"], FileTransferMethod.LOCAL_FILE),
("text/plain", b"Hello, world!", ["Hello, world!"], FileTransferMethod.LOCAL_FILE, ".txt"),
(
"application/pdf",
b"%PDF-1.5\n%Test PDF content",
["Mocked PDF content"],
FileTransferMethod.LOCAL_FILE,
".pdf",
),
(
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
b"PK\x03\x04",
["Mocked DOCX content"],
FileTransferMethod.LOCAL_FILE,
FileTransferMethod.REMOTE_URL,
"",
),
("text/plain", b"Remote content", ["Remote content"], FileTransferMethod.REMOTE_URL),
("text/plain", b"Remote content", ["Remote content"], FileTransferMethod.REMOTE_URL, None),
],
)
def test_run_extract_text(
@ -83,6 +90,7 @@ def test_run_extract_text(
file_content,
expected_text,
transfer_method,
extension,
monkeypatch,
):
document_extractor_node.graph_runtime_state = mock_graph_runtime_state
@ -92,6 +100,7 @@ def test_run_extract_text(
mock_file.transfer_method = transfer_method
mock_file.related_id = "test_file_id" if transfer_method == FileTransferMethod.LOCAL_FILE else None
mock_file.remote_url = "https://example.com/file.txt" if transfer_method == FileTransferMethod.REMOTE_URL else None
mock_file.extension = extension
mock_array_file_segment = Mock(spec=ArrayFileSegment)
mock_array_file_segment.value = [mock_file]
@ -116,7 +125,7 @@ def test_run_extract_text(
result = document_extractor_node._run()
assert isinstance(result, NodeRunResult)
assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED, result.error
assert result.outputs is not None
assert result.outputs["text"] == expected_text

View File

@ -192,7 +192,7 @@ def test_http_request_node_form_with_file(monkeypatch):
def attr_checker(*args, **kwargs):
assert kwargs["data"] == {"name": "test"}
assert kwargs["files"] == {"file": b"test"}
assert kwargs["files"] == {"file": (None, b"test", "application/octet-stream")}
return httpx.Response(200, content=b"")
monkeypatch.setattr(

View File

@ -55,6 +55,7 @@ def test_execute_if_else_result_true():
pool.add(["start", "less_than"], 21)
pool.add(["start", "greater_than_or_equal"], 22)
pool.add(["start", "less_than_or_equal"], 21)
pool.add(["start", "null"], None)
pool.add(["start", "not_null"], "1212")
node = IfElseNode(

View File

@ -33,8 +33,8 @@ def test_get_file_attribute(pool, file):
assert result.value == file.filename
# Test getting a non-existent attribute
with pytest.raises(ValueError):
pool.get(("node_1", "file_var", "non_existent_attr"))
result = pool.get(("node_1", "file_var", "non_existent_attr"))
assert result is None
def test_use_long_selector(pool):

View File

@ -0,0 +1,100 @@
import os
import posixpath
from unittest.mock import MagicMock
import pytest
from _pytest.monkeypatch import MonkeyPatch
from oss2 import Bucket
from oss2.models import GetObjectResult, PutObjectResult
from tests.unit_tests.oss.__mock.base import (
get_example_bucket,
get_example_data,
get_example_filename,
get_example_filepath,
get_example_folder,
)
class MockResponse:
def __init__(self, status, headers, request_id):
self.status = status
self.headers = headers
self.request_id = request_id
class MockAliyunOssClass:
def __init__(
self,
auth,
endpoint,
bucket_name,
is_cname=False,
session=None,
connect_timeout=None,
app_name="",
enable_crc=True,
proxies=None,
region=None,
cloudbox_id=None,
is_path_style=False,
is_verify_object_strict=True,
):
self.bucket_name = get_example_bucket()
self.key = posixpath.join(get_example_folder(), get_example_filename())
self.content = get_example_data()
self.filepath = get_example_filepath()
self.resp = MockResponse(
200,
{
"etag": "ee8de918d05640145b18f70f4c3aa602",
"x-oss-version-id": "CAEQNhiBgMDJgZCA0BYiIDc4MGZjZGI2OTBjOTRmNTE5NmU5NmFhZjhjYmY0****",
},
"request_id",
)
def put_object(self, key, data, headers=None, progress_callback=None):
assert key == self.key
assert data == self.content
return PutObjectResult(self.resp)
def get_object(self, key, byte_range=None, headers=None, progress_callback=None, process=None, params=None):
assert key == self.key
get_object_output = MagicMock(GetObjectResult)
get_object_output.read.return_value = self.content
return get_object_output
def get_object_to_file(
self, key, filename, byte_range=None, headers=None, progress_callback=None, process=None, params=None
):
assert key == self.key
assert filename == self.filepath
def object_exists(self, key, headers=None):
assert key == self.key
return True
def delete_object(self, key, params=None, headers=None):
assert key == self.key
self.resp.headers["x-oss-delete-marker"] = True
return self.resp
MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
@pytest.fixture
def setup_aliyun_oss_mock(monkeypatch: MonkeyPatch):
if MOCK:
monkeypatch.setattr(Bucket, "__init__", MockAliyunOssClass.__init__)
monkeypatch.setattr(Bucket, "put_object", MockAliyunOssClass.put_object)
monkeypatch.setattr(Bucket, "get_object", MockAliyunOssClass.get_object)
monkeypatch.setattr(Bucket, "get_object_to_file", MockAliyunOssClass.get_object_to_file)
monkeypatch.setattr(Bucket, "object_exists", MockAliyunOssClass.object_exists)
monkeypatch.setattr(Bucket, "delete_object", MockAliyunOssClass.delete_object)
yield
if MOCK:
monkeypatch.undo()

View File

@ -0,0 +1,58 @@
from collections.abc import Generator
import pytest
from extensions.storage.base_storage import BaseStorage
def get_example_folder() -> str:
return "/dify"
def get_example_bucket() -> str:
return "dify"
def get_example_filename() -> str:
return "test.txt"
def get_example_data() -> bytes:
return b"test"
def get_example_filepath() -> str:
return "/test"
class BaseStorageTest:
@pytest.fixture(autouse=True)
def setup_method(self):
"""Should be implemented in child classes to setup specific storage."""
self.storage = BaseStorage()
def test_save(self):
"""Test saving data."""
self.storage.save(get_example_filename(), get_example_data())
def test_load_once(self):
"""Test loading data once."""
assert self.storage.load_once(get_example_filename()) == get_example_data()
def test_load_stream(self):
"""Test loading data as a stream."""
generator = self.storage.load_stream(get_example_filename())
assert isinstance(generator, Generator)
assert next(generator) == get_example_data()
def test_download(self):
"""Test downloading data."""
self.storage.download(get_example_filename(), get_example_filepath())
def test_exists(self):
"""Test checking if a file exists."""
assert self.storage.exists(get_example_filename())
def test_delete(self):
"""Test deleting a file."""
self.storage.delete(get_example_filename())

View File

@ -0,0 +1,57 @@
import os
import shutil
from pathlib import Path
from unittest.mock import MagicMock, mock_open, patch
import pytest
from _pytest.monkeypatch import MonkeyPatch
from tests.unit_tests.oss.__mock.base import (
get_example_data,
get_example_filename,
get_example_filepath,
get_example_folder,
)
class MockLocalFSClass:
def write_bytes(self, data):
assert data == get_example_data()
def read_bytes(self):
return get_example_data()
@staticmethod
def copyfile(src, dst):
assert src == os.path.join(get_example_folder(), get_example_filename())
assert dst == get_example_filepath()
@staticmethod
def exists(path):
assert path == os.path.join(get_example_folder(), get_example_filename())
return True
@staticmethod
def remove(path):
assert path == os.path.join(get_example_folder(), get_example_filename())
MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
@pytest.fixture
def setup_local_fs_mock(monkeypatch: MonkeyPatch):
if MOCK:
monkeypatch.setattr(Path, "write_bytes", MockLocalFSClass.write_bytes)
monkeypatch.setattr(Path, "read_bytes", MockLocalFSClass.read_bytes)
monkeypatch.setattr(shutil, "copyfile", MockLocalFSClass.copyfile)
monkeypatch.setattr(os.path, "exists", MockLocalFSClass.exists)
monkeypatch.setattr(os, "remove", MockLocalFSClass.remove)
os.makedirs = MagicMock()
with patch("builtins.open", mock_open(read_data=get_example_data())):
yield
if MOCK:
monkeypatch.undo()

View File

@ -0,0 +1,81 @@
import os
from unittest.mock import MagicMock
import pytest
from _pytest.monkeypatch import MonkeyPatch
from qcloud_cos import CosS3Client
from qcloud_cos.streambody import StreamBody
from tests.unit_tests.oss.__mock.base import (
get_example_bucket,
get_example_data,
get_example_filename,
get_example_filepath,
)
class MockTencentCosClass:
def __init__(self, conf, retry=1, session=None):
self.bucket_name = get_example_bucket()
self.key = get_example_filename()
self.content = get_example_data()
self.filepath = get_example_filepath()
self.resp = {
"ETag": "ee8de918d05640145b18f70f4c3aa602",
"Server": "tencent-cos",
"x-cos-hash-crc64ecma": 16749565679157681890,
"x-cos-request-id": "NWU5MDNkYzlfNjRiODJhMDlfMzFmYzhfMTFm****",
}
def put_object(self, Bucket, Body, Key, EnableMD5=False, **kwargs): # noqa: N803
assert Bucket == self.bucket_name
assert Key == self.key
assert Body == self.content
return self.resp
def get_object(self, Bucket, Key, KeySimplifyCheck=True, **kwargs): # noqa: N803
assert Bucket == self.bucket_name
assert Key == self.key
mock_stream_body = MagicMock(StreamBody)
mock_raw_stream = MagicMock()
mock_stream_body.get_raw_stream.return_value = mock_raw_stream
mock_raw_stream.read.return_value = self.content
mock_stream_body.get_stream_to_file = MagicMock()
def chunk_generator(chunk_size=2):
for i in range(0, len(self.content), chunk_size):
yield self.content[i : i + chunk_size]
mock_stream_body.get_stream.return_value = chunk_generator(chunk_size=4096)
return {"Body": mock_stream_body}
def object_exists(self, Bucket, Key): # noqa: N803
assert Bucket == self.bucket_name
assert Key == self.key
return True
def delete_object(self, Bucket, Key, **kwargs): # noqa: N803
assert Bucket == self.bucket_name
assert Key == self.key
self.resp.update({"x-cos-delete-marker": True})
return self.resp
MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
@pytest.fixture
def setup_tencent_cos_mock(monkeypatch: MonkeyPatch):
if MOCK:
monkeypatch.setattr(CosS3Client, "__init__", MockTencentCosClass.__init__)
monkeypatch.setattr(CosS3Client, "put_object", MockTencentCosClass.put_object)
monkeypatch.setattr(CosS3Client, "get_object", MockTencentCosClass.get_object)
monkeypatch.setattr(CosS3Client, "object_exists", MockTencentCosClass.object_exists)
monkeypatch.setattr(CosS3Client, "delete_object", MockTencentCosClass.delete_object)
yield
if MOCK:
monkeypatch.undo()

View File

@ -1,5 +1,4 @@
import os
from typing import Union
from unittest.mock import MagicMock
import pytest
@ -7,28 +6,19 @@ from _pytest.monkeypatch import MonkeyPatch
from tos import TosClientV2
from tos.clientv2 import DeleteObjectOutput, GetObjectOutput, HeadObjectOutput, PutObjectOutput
from tests.unit_tests.oss.__mock.base import (
get_example_bucket,
get_example_data,
get_example_filename,
get_example_filepath,
)
class AttrDict(dict):
def __getattr__(self, item):
return self.get(item)
def get_example_bucket() -> str:
return "dify"
def get_example_filename() -> str:
return "test.txt"
def get_example_data() -> bytes:
return b"test"
def get_example_filepath() -> str:
return "/test"
class MockVolcengineTosClass:
def __init__(self, ak="", sk="", endpoint="", region=""):
self.bucket_name = get_example_bucket()

View File

@ -0,0 +1,22 @@
from unittest.mock import MagicMock, patch
import pytest
from oss2 import Auth
from extensions.storage.aliyun_oss_storage import AliyunOssStorage
from tests.unit_tests.oss.__mock.aliyun_oss import setup_aliyun_oss_mock
from tests.unit_tests.oss.__mock.base import (
BaseStorageTest,
get_example_bucket,
get_example_folder,
)
class TestAliyunOss(BaseStorageTest):
@pytest.fixture(autouse=True)
def setup_method(self, setup_aliyun_oss_mock):
"""Executed before each test method."""
with patch.object(Auth, "__init__", return_value=None):
self.storage = AliyunOssStorage()
self.storage.bucket_name = get_example_bucket()
self.storage.folder = get_example_folder()

View File

@ -0,0 +1,18 @@
from collections.abc import Generator
import pytest
from extensions.storage.local_fs_storage import LocalFsStorage
from tests.unit_tests.oss.__mock.base import (
BaseStorageTest,
get_example_folder,
)
from tests.unit_tests.oss.__mock.local import setup_local_fs_mock
class TestLocalFS(BaseStorageTest):
@pytest.fixture(autouse=True)
def setup_method(self, setup_local_fs_mock):
"""Executed before each test method."""
self.storage = LocalFsStorage()
self.storage.folder = get_example_folder()

View File

@ -0,0 +1,20 @@
from unittest.mock import patch
import pytest
from qcloud_cos import CosConfig
from extensions.storage.tencent_cos_storage import TencentCosStorage
from tests.unit_tests.oss.__mock.base import (
BaseStorageTest,
get_example_bucket,
)
from tests.unit_tests.oss.__mock.tencent_cos import setup_tencent_cos_mock
class TestTencentCos(BaseStorageTest):
@pytest.fixture(autouse=True)
def setup_method(self, setup_tencent_cos_mock):
"""Executed before each test method."""
with patch.object(CosConfig, "__init__", return_value=None):
self.storage = TencentCosStorage()
self.storage.bucket_name = get_example_bucket()

View File

@ -1,30 +1,18 @@
from collections.abc import Generator
from flask import Flask
import pytest
from tos import TosClientV2
from tos.clientv2 import GetObjectOutput, HeadObjectOutput, PutObjectOutput
from extensions.storage.volcengine_tos_storage import VolcengineTosStorage
from tests.unit_tests.oss.__mock.volcengine_tos import (
from tests.unit_tests.oss.__mock.base import (
BaseStorageTest,
get_example_bucket,
get_example_data,
get_example_filename,
get_example_filepath,
setup_volcengine_tos_mock,
)
from tests.unit_tests.oss.__mock.volcengine_tos import setup_volcengine_tos_mock
class VolcengineTosTest:
_instance = None
def __new__(cls):
if cls._instance == None:
cls._instance = object.__new__(cls)
return cls._instance
else:
return cls._instance
def __init__(self):
class TestVolcengineTos(BaseStorageTest):
@pytest.fixture(autouse=True)
def setup_method(self, setup_volcengine_tos_mock):
"""Executed before each test method."""
self.storage = VolcengineTosStorage()
self.storage.bucket_name = get_example_bucket()
self.storage.client = TosClientV2(
@ -33,35 +21,3 @@ class VolcengineTosTest:
endpoint="https://xxx.volces.com",
region="cn-beijing",
)
def test_save(setup_volcengine_tos_mock):
volc_tos = VolcengineTosTest()
volc_tos.storage.save(get_example_filename(), get_example_data())
def test_load_once(setup_volcengine_tos_mock):
volc_tos = VolcengineTosTest()
assert volc_tos.storage.load_once(get_example_filename()) == get_example_data()
def test_load_stream(setup_volcengine_tos_mock):
volc_tos = VolcengineTosTest()
generator = volc_tos.storage.load_stream(get_example_filename())
assert isinstance(generator, Generator)
assert next(generator) == get_example_data()
def test_download(setup_volcengine_tos_mock):
volc_tos = VolcengineTosTest()
volc_tos.storage.download(get_example_filename(), get_example_filepath())
def test_exists(setup_volcengine_tos_mock):
volc_tos = VolcengineTosTest()
assert volc_tos.storage.exists(get_example_filename())
def test_delete(setup_volcengine_tos_mock):
volc_tos = VolcengineTosTest()
volc_tos.storage.delete(get_example_filename())

View File

@ -0,0 +1,47 @@
import pytest
from packaging import version
from services.app_dsl_service import AppDslService
from services.app_dsl_service.exc import DSLVersionNotSupportedError
from services.app_dsl_service.service import _check_or_fix_dsl, current_dsl_version
class TestAppDSLService:
@pytest.mark.skip(reason="Test skipped")
def test_check_or_fix_dsl_missing_version(self):
import_data = {}
result = _check_or_fix_dsl(import_data)
assert result["version"] == "0.1.0"
assert result["kind"] == "app"
@pytest.mark.skip(reason="Test skipped")
def test_check_or_fix_dsl_missing_kind(self):
import_data = {"version": "0.1.0"}
result = _check_or_fix_dsl(import_data)
assert result["kind"] == "app"
@pytest.mark.skip(reason="Test skipped")
def test_check_or_fix_dsl_older_version(self):
import_data = {"version": "0.0.9", "kind": "app"}
result = _check_or_fix_dsl(import_data)
assert result["version"] == "0.0.9"
@pytest.mark.skip(reason="Test skipped")
def test_check_or_fix_dsl_current_version(self):
import_data = {"version": current_dsl_version, "kind": "app"}
result = _check_or_fix_dsl(import_data)
assert result["version"] == current_dsl_version
@pytest.mark.skip(reason="Test skipped")
def test_check_or_fix_dsl_newer_version(self):
current_version = version.parse(current_dsl_version)
newer_version = f"{current_version.major}.{current_version.minor + 1}.0"
import_data = {"version": newer_version, "kind": "app"}
with pytest.raises(DSLVersionNotSupportedError):
_check_or_fix_dsl(import_data)
@pytest.mark.skip(reason="Test skipped")
def test_check_or_fix_dsl_invalid_kind(self):
import_data = {"version": current_dsl_version, "kind": "invalid"}
result = _check_or_fix_dsl(import_data)
assert result["kind"] == "app"