github-mirrors/dify
1
0
Fork 0
mirror of https://mirror.skon.top/github.com/langgenius/dify.git synced 2026-05-01 03:30:02 +08:00
Code Issues Actions 7 Packages Projects Releases Wiki Activity

fix unit test

Browse Source
This commit is contained in:
Yansong Zhang
2026-04-15 12:19:40 +08:00
parent c08986ad78
commit fa9aaa91fd
1 changed files with 34 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
api/providers/vdb/vdb-tidb-on-qdrant/tests/unit_tests/test_tidb_service.py
View File

@@ -3,13 +3,33 @@ from unittest.mock import MagicMock, patch
from dify_vdb_tidb_on_qdrant.tidb_service import TidbService
class TestExtractQdrantEndpoint:
"""Unit tests for TidbService.extract_qdrant_endpoint."""
def test_returns_endpoint_when_host_present(self):
response = {"endpoints": {"public": {"host": "gateway01.us-east-1.tidbcloud.com", "port": 4000}}}
result = TidbService.extract_qdrant_endpoint(response)
assert result == "https://qdrant-gateway01.us-east-1.tidbcloud.com"
def test_returns_none_when_host_missing(self):
response = {"endpoints": {"public": {}}}
assert TidbService.extract_qdrant_endpoint(response) is None
def test_returns_none_when_public_missing(self):
response = {"endpoints": {}}
assert TidbService.extract_qdrant_endpoint(response) is None
def test_returns_none_when_endpoints_missing(self):
assert TidbService.extract_qdrant_endpoint({}) is None
class TestFetchQdrantEndpoint:
"""Unit tests for TidbService.fetch_qdrant_endpoint."""
@patch.object(TidbService, "get_tidb_serverless_cluster")
def test_returns_endpoint_when_host_present(self, mock_get_cluster):
mock_get_cluster.return_value = {
"status": {"connection_strings": {"standard": {"host": "gateway01.us-east-1.tidbcloud.com"}}}
"endpoints": {"public": {"host": "gateway01.us-east-1.tidbcloud.com", "port": 4000}}
}
result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
assert result == "https://qdrant-gateway01.us-east-1.tidbcloud.com"
@@ -17,65 +37,48 @@ class TestFetchQdrantEndpoint:
@patch.object(TidbService, "get_tidb_serverless_cluster")
def test_returns_none_when_cluster_response_is_none(self, mock_get_cluster):
mock_get_cluster.return_value = None
result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
assert result is None
assert TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123") is None
@patch.object(TidbService, "get_tidb_serverless_cluster")
def test_returns_none_when_host_missing(self, mock_get_cluster):
mock_get_cluster.return_value = {"status": {"connection_strings": {"standard": {}}}}
result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
assert result is None
mock_get_cluster.return_value = {"endpoints": {"public": {}}}
assert TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123") is None
@patch.object(TidbService, "get_tidb_serverless_cluster")
def test_returns_none_when_status_missing(self, mock_get_cluster):
def test_returns_none_when_endpoints_missing(self, mock_get_cluster):
mock_get_cluster.return_value = {}
result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
assert result is None
@patch.object(TidbService, "get_tidb_serverless_cluster")
def test_returns_none_when_connection_strings_missing(self, mock_get_cluster):
mock_get_cluster.return_value = {"status": {}}
result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
assert result is None
assert TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123") is None
@patch.object(TidbService, "get_tidb_serverless_cluster")
def test_returns_none_on_exception(self, mock_get_cluster):
mock_get_cluster.side_effect = RuntimeError("network error")
result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
assert result is None
@patch.object(TidbService, "get_tidb_serverless_cluster")
def test_returns_none_when_standard_key_missing(self, mock_get_cluster):
mock_get_cluster.return_value = {"status": {"connection_strings": {}}}
result = TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123")
assert result is None
assert TidbService.fetch_qdrant_endpoint("url", "pub", "priv", "c-123") is None
class TestCreateTidbServerlessClusterQdrantEndpoint:
"""Verify that create_tidb_serverless_cluster includes qdrant_endpoint in its result."""
@patch.object(TidbService, "fetch_qdrant_endpoint", return_value="https://qdrant-gw.tidbcloud.com")
@patch.object(TidbService, "get_tidb_serverless_cluster")
@patch("dify_vdb_tidb_on_qdrant.tidb_service._tidb_http_client")
@patch("dify_vdb_tidb_on_qdrant.tidb_service.dify_config")
def test_result_contains_qdrant_endpoint(self, mock_config, mock_http, mock_get_cluster, mock_fetch_ep):
def test_result_contains_qdrant_endpoint(self, mock_config, mock_http, mock_get_cluster):
mock_config.TIDB_SPEND_LIMIT = 10
mock_http.post.return_value = MagicMock(status_code=200, json=lambda: {"clusterId": "c-1"})
mock_get_cluster.return_value = {"state": "ACTIVE", "userPrefix": "pfx"}
mock_get_cluster.return_value = {
"state": "ACTIVE",
"userPrefix": "pfx",
"endpoints": {"public": {"host": "gw.tidbcloud.com", "port": 4000}},
}
result = TidbService.create_tidb_serverless_cluster("proj", "url", "iam", "pub", "priv", "us-east-1")
assert result is not None
assert result["qdrant_endpoint"] == "https://qdrant-gw.tidbcloud.com"
mock_fetch_ep.assert_called_once_with("url", "pub", "priv", "c-1")
@patch.object(TidbService, "fetch_qdrant_endpoint", return_value=None)
@patch.object(TidbService, "get_tidb_serverless_cluster")
@patch("dify_vdb_tidb_on_qdrant.tidb_service._tidb_http_client")
@patch("dify_vdb_tidb_on_qdrant.tidb_service.dify_config")
def test_result_qdrant_endpoint_none_when_fetch_fails(
self, mock_config, mock_http, mock_get_cluster, mock_fetch_ep
):
def test_result_qdrant_endpoint_none_when_no_endpoints(self, mock_config, mock_http, mock_get_cluster):
mock_config.TIDB_SPEND_LIMIT = 10
mock_http.post.return_value = MagicMock(status_code=200, json=lambda: {"clusterId": "c-1"})
mock_get_cluster.return_value = {"state": "ACTIVE", "userPrefix": "pfx"}