diff --git a/tools/tianditu/manifest.yaml b/tools/tianditu/manifest.yaml index 903db0308..78f781699 100644 --- a/tools/tianditu/manifest.yaml +++ b/tools/tianditu/manifest.yaml @@ -36,4 +36,4 @@ tags: - utilities - travel type: plugin -version: 0.0.7 +version: 0.0.8 diff --git a/tools/tianditu/tests/__init__.py b/tools/tianditu/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tools/tianditu/tests/test_tianditu.py b/tools/tianditu/tests/test_tianditu.py new file mode 100644 index 000000000..57ae77feb --- /dev/null +++ b/tools/tianditu/tests/test_tianditu.py @@ -0,0 +1,316 @@ +"""In-process tests for tools/tianditu (no pytest required). + +These tests cover two reliability fixes shipped in this PR: + +- Bug 1: missing `timeout=` on `requests.get` calls. +- Bug 2: `KeyError` when `tianditu_api_key` is missing from credentials. + +`tianditu` does not take user-supplied `hl` / `gl` (the geocoder API takes a +`keyWord` and an optional `region` already from upstream, not through us), +so the third bug from PR #3456 (invoke fall-through on invalid hl/gl) +does not apply. + +They use only Python stdlib (`unittest.mock`) and tiny in-test stubs for +`requests`, `dify_plugin`, and `pydantic` so they can run in environments +where the plugin dependencies are not pre-installed. +""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +_HERE = Path(__file__).resolve().parent # tests/ +_PLUGIN_ROOT = _HERE.parent # tools/tianditu/ +sys.path.insert(0, str(_PLUGIN_ROOT)) # so `tools.geocoder` etc. resolve + + +# ---- Stub modules ---------------------------------------------------------- + +def _ensure_stub_modules() -> None: + if "requests" not in sys.modules: + requests_stub = types.ModuleType("requests") + requests_stub.exceptions = types.SimpleNamespace( + RequestException=type("RequestException", (Exception,), {}) + ) + + def _default_get(*args, **kwargs): + raise AssertionError( + "requests.get was called in a test that did not patch it" + ) + + requests_stub.get = _default_get + sys.modules["requests"] = requests_stub + + if "dify_plugin" not in sys.modules: + dify_plugin_stub = types.ModuleType("dify_plugin") + + class _BaseTool: + def __init__(self): + self.runtime = SimpleNamespace(credentials={}) + + def create_text_message(self, text): + return ("text", text) + + def create_json_message(self, json): + return ("json", json) + + def create_blob_message(self, blob, meta=None): + return ("blob", blob, meta) + + dify_plugin_stub.Tool = _BaseTool + + dify_plugin_entities_stub = types.ModuleType("dify_plugin.entities") + tool_module_stub = types.ModuleType("dify_plugin.entities.tool") + + class _ToolInvokeMessage: + def __init__(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + + tool_module_stub.ToolInvokeMessage = _ToolInvokeMessage + + dify_plugin_entities_stub.tool = tool_module_stub + sys.modules["dify_plugin"] = dify_plugin_stub + sys.modules["dify_plugin.entities"] = dify_plugin_entities_stub + sys.modules["dify_plugin.entities.tool"] = tool_module_stub + + +_ensure_stub_modules() + +# Force-fresh imports so we get the patched stubs above. +for _name in list(sys.modules): + if _name.startswith("tools.tianditu") or _name.startswith("tools.geocoder") \ + or _name.startswith("tools.staticmap") or _name.startswith("tools.poisearch"): + sys.modules.pop(_name, None) + +from importlib import util as _importlib_util # noqa: E402 + + +def _load_module(file_name: str): + path = _PLUGIN_ROOT / "tools" / file_name + spec = _importlib_util.spec_from_file_location(f"tianditu.{file_name[:-3]}", path) + mod = _importlib_util.module_from_spec(spec) + sys.modules[spec.name] = mod + spec.loader.exec_module(mod) + return mod + + +geocoder = _load_module("geocoder.py") +staticmap = _load_module("staticmap.py") +poisearch = _load_module("poisearch.py") + + +# ---- Helpers --------------------------------------------------------------- + +class _FakeResponse: + def __init__(self, payload=None, *, raw_bytes=b"\x89PNG\r\n", status_code=200): + if payload is not None: + self._payload = payload + else: + self._payload = { + "status": "0", + "message": "ok", + "location": {"lon": 116.404, "lat": 39.915}, + } + self._raw_bytes = raw_bytes + self.status_code = status_code + self.calls = 0 + + def raise_for_status(self): + if self.status_code >= 400: + raise geocoder.requests.exceptions.RequestException( + f"HTTP {self.status_code}" + ) + + def json(self): + self.calls += 1 + return self._payload + + @property + def content(self): + return self._raw_bytes + + +def _make_tool(tool_cls, credentials): + tool = object.__new__(tool_cls) + + def _text(text): + return ("text", text) + + def _json(json=None, **_kw): + return ("json", json) + + def _blob(blob, meta=None): + return ("blob", blob, meta) + + tool.create_text_message = _text + tool.create_json_message = _json + tool.create_blob_message = _blob + tool.runtime = SimpleNamespace(credentials=credentials) + return tool + + +def _flatten(messages): + return list(messages) + + +# ============================================================================= +# Bug 2: missing-credential guard +# ============================================================================= + +def test_geocoder_missing_api_key_returns_message_and_no_http_call(): + tool = _make_tool(geocoder.GeocoderTool, credentials={}) + fake_get = mock.Mock(side_effect=AssertionError("requests.get must not be called")) + with mock.patch.object(geocoder.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert len(messages) == 1 + assert messages[0][0] == "text" + assert messages[0][1] == "Tianditu API key is required." + assert fake_get.call_count == 0 + + +def test_geocoder_empty_api_key_returns_message_and_no_http_call(): + tool = _make_tool(geocoder.GeocoderTool, credentials={"tianditu_api_key": ""}) + fake_get = mock.Mock(side_effect=AssertionError("requests.get must not be called")) + with mock.patch.object(geocoder.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert messages[0][1] == "Tianditu API key is required." + assert fake_get.call_count == 0 + + +def test_staticmap_missing_api_key_returns_message_and_no_http_call(): + tool = _make_tool(staticmap.PoiSearchTool, credentials={}) + fake_get = mock.Mock(side_effect=AssertionError("requests.get must not be called")) + with mock.patch.object(staticmap.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert messages[0][1] == "Tianditu API key is required." + assert fake_get.call_count == 0 + + +def test_poisearch_missing_api_key_returns_message_and_no_http_call(): + tool = _make_tool(poisearch.PoiSearchTool, credentials={}) + fake_get = mock.Mock(side_effect=AssertionError("requests.get must not be called")) + with mock.patch.object(poisearch.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing", "baseAddress": "Beijing"})) + assert messages[0][1] == "Tianditu API key is required." + assert fake_get.call_count == 0 + + +# ============================================================================= +# Bug 2 corollary: missing key did NOT raise KeyError pre-fix. +# ============================================================================= + +def test_geocoder_missing_api_key_does_not_raise_keyerror(): + tool = _make_tool(geocoder.GeocoderTool, credentials={}) + fake_get = mock.Mock(side_effect=AssertionError("requests.get must not be called")) + with mock.patch.object(geocoder.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert messages[0] == ("text", "Tianditu API key is required.") + + +# ============================================================================= +# Bug 1: requests.get forwards a 10s timeout on every call. +# ============================================================================= + +def test_geocoder_requests_get_uses_10s_timeout(): + tool = _make_tool(geocoder.GeocoderTool, credentials={"tianditu_api_key": "secret"}) + fake_response = _FakeResponse(payload={"status": "0", "message": "ok"}) + fake_get = mock.Mock(return_value=fake_response) + with mock.patch.object(geocoder.requests, "get", fake_get): + _flatten(tool._invoke({"keyword": "Beijing"})) + assert fake_get.call_count == 1 + _, kwargs = fake_get.call_args + assert kwargs.get("timeout") == 10, f"expected timeout=10 kwarg, got {kwargs!r}" + + +def test_staticmap_requests_get_uses_10s_timeout_on_every_call(): + tool = _make_tool(staticmap.PoiSearchTool, credentials={"tianditu_api_key": "secret"}) + fake_response = _FakeResponse() + fake_get = mock.Mock(return_value=fake_response) + with mock.patch.object(staticmap.requests, "get", fake_get): + _flatten(tool._invoke({"keyword": "Beijing"})) + # staticmap makes 2 calls: geocoder look-up + static image fetch. + assert fake_get.call_count == 2 + for call in fake_get.call_args_list: + _, kwargs = call + assert kwargs.get("timeout") == 10, f"expected timeout=10 kwarg, got {kwargs!r}" + + +def test_poisearch_requests_get_uses_10s_timeout_on_every_call(): + tool = _make_tool(poisearch.PoiSearchTool, credentials={"tianditu_api_key": "secret"}) + fake_response = _FakeResponse() + fake_get = mock.Mock(return_value=fake_response) + with mock.patch.object(poisearch.requests, "get", fake_get): + _flatten(tool._invoke({"keyword": "Beijing", "baseAddress": "Beijing"})) + # poisearch makes 2 calls: geocoder look-up + v2/search + assert fake_get.call_count == 2 + for call in fake_get.call_args_list: + _, kwargs = call + assert kwargs.get("timeout") == 10, f"expected timeout=10 kwarg, got {kwargs!r}" + + +# ============================================================================= +# Happy-path smoke +# ============================================================================= + +def test_geocoder_happy_path_yields_json_message(): + tool = _make_tool(geocoder.GeocoderTool, credentials={"tianditu_api_key": "secret"}) + fake_response = _FakeResponse() + fake_get = mock.Mock(return_value=fake_response) + with mock.patch.object(geocoder.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert any(m[0] == "json" for m in messages) + + +def test_staticmap_happy_path_yields_blob_message(): + tool = _make_tool(staticmap.PoiSearchTool, credentials={"tianditu_api_key": "secret"}) + fake_response = _FakeResponse() + fake_get = mock.Mock(return_value=fake_response) + with mock.patch.object(staticmap.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert any(m[0] == "blob" for m in messages) + + +# ============================================================================= +# Network failure surfaces as a friendly message, not a stack trace. +# ============================================================================= + +def test_geocoder_request_exception_yields_friendly_message(): + tool = _make_tool(geocoder.GeocoderTool, credentials={"tianditu_api_key": "secret"}) + fake_get = mock.Mock(side_effect=geocoder.requests.exceptions.RequestException("boom")) + with mock.patch.object(geocoder.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert any( + m[0] == "text" and "network or upstream failure" in m[1] for m in messages + ), f"expected generic friendly error, got {messages!r}" + assert not any( + m[0] == "text" and "boom" in m[1] for m in messages + ), "exception text must not be forwarded to the user" + + +def test_geocoder_request_does_not_leak_tk_in_message(): + tool = _make_tool(geocoder.GeocoderTool, credentials={"tianditu_api_key": "secret-key"}) + fake_get = mock.Mock( + side_effect=geocoder.requests.exceptions.RequestException( + "GET http://api.tianditu.gov.cn/geocoder?ds=%7B%22keyWord%22%3A%22Beijing%22%7D&tk=secret-key" + ) + ) + with mock.patch.object(geocoder.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + for m in messages: + if m[0] == "text": + assert "secret-key" not in m[1], f"API key leaked in message: {m[1]!r}" + + +def test_geocoder_http_4xx_returns_message_and_no_json_dump(): + tool = _make_tool(geocoder.GeocoderTool, credentials={"tianditu_api_key": "secret"}) + fake_response = _FakeResponse(status_code=500) + fake_get = mock.Mock(return_value=fake_response) + with mock.patch.object(geocoder.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert any(m[0] == "text" and "network or upstream failure" in m[1] for m in messages) + assert not any(m[0] == "json" for m in messages) diff --git a/tools/tianditu/tools/geocoder.py b/tools/tianditu/tools/geocoder.py index fc07e0e52..a11b1d4b7 100644 --- a/tools/tianditu/tools/geocoder.py +++ b/tools/tianditu/tools/geocoder.py @@ -4,6 +4,11 @@ from dify_plugin.entities.tool import ToolInvokeMessage from dify_plugin import Tool +_NETWORK_ERROR = ( + "An error occurred while invoking the tool (network or upstream failure). " + "Please retry shortly." +) + class GeocoderTool(Tool): def _invoke( @@ -17,9 +22,21 @@ def _invoke( if not keyword: yield self.create_text_message("Invalid parameter keyword") return - tk = self.runtime.credentials["tianditu_api_key"] + + tk = self.runtime.credentials.get("tianditu_api_key") + if not tk: + yield self.create_text_message("Tianditu API key is required.") + return + params = {"keyWord": keyword} - result = requests.get( - base_url + "?ds=" + json.dumps(params, ensure_ascii=False) + "&tk=" + tk - ).json() + try: + response = requests.get( + base_url + "?ds=" + json.dumps(params, ensure_ascii=False) + "&tk=" + tk, + timeout=10, + ) + response.raise_for_status() + result = response.json() + except requests.exceptions.RequestException: + yield self.create_text_message(_NETWORK_ERROR) + return yield self.create_json_message(result) diff --git a/tools/tianditu/tools/poisearch.py b/tools/tianditu/tools/poisearch.py index efd08aa7f..325231af0 100644 --- a/tools/tianditu/tools/poisearch.py +++ b/tools/tianditu/tools/poisearch.py @@ -4,6 +4,11 @@ from dify_plugin.entities.tool import ToolInvokeMessage from dify_plugin import Tool +_NETWORK_ERROR = ( + "An error occurred while invoking the tool (network or upstream failure). " + "Please retry shortly." +) + class PoiSearchTool(Tool): def _invoke( @@ -22,29 +27,44 @@ def _invoke( if not baseAddress: yield self.create_text_message("Invalid parameter baseAddress") return - tk = self.runtime.credentials["tianditu_api_key"] - base_coords = requests.get( - geocoder_base_url - + "?ds=" - + json.dumps({"keyWord": baseAddress}, ensure_ascii=False) - + "&tk=" - + tk - ).json() - params = { - "keyWord": keyword, - "queryRadius": 5000, - "queryType": 3, - "pointLonlat": str(base_coords["location"]["lon"]) - + "," - + str(base_coords["location"]["lat"]), - "start": 0, - "count": 100, - } - result = requests.get( - base_url - + "?postStr=" - + json.dumps(params, ensure_ascii=False) - + "&type=query&tk=" - + tk - ).json() + + tk = self.runtime.credentials.get("tianditu_api_key") + if not tk: + yield self.create_text_message("Tianditu API key is required.") + return + + try: + geo_response = requests.get( + geocoder_base_url + + "?ds=" + + json.dumps({"keyWord": baseAddress}, ensure_ascii=False) + + "&tk=" + + tk, + timeout=10, + ) + geo_response.raise_for_status() + base_coords = geo_response.json() + params = { + "keyWord": keyword, + "queryRadius": 5000, + "queryType": 3, + "pointLonlat": str(base_coords["location"]["lon"]) + + "," + + str(base_coords["location"]["lat"]), + "start": 0, + "count": 100, + } + search_response = requests.get( + base_url + + "?postStr=" + + json.dumps(params, ensure_ascii=False) + + "&type=query&tk=" + + tk, + timeout=10, + ) + search_response.raise_for_status() + result = search_response.json() + except requests.exceptions.RequestException: + yield self.create_text_message(_NETWORK_ERROR) + return yield self.create_json_message(result) diff --git a/tools/tianditu/tools/staticmap.py b/tools/tianditu/tools/staticmap.py index 00a1e83f6..af1844866 100644 --- a/tools/tianditu/tools/staticmap.py +++ b/tools/tianditu/tools/staticmap.py @@ -4,6 +4,11 @@ from dify_plugin.entities.tool import ToolInvokeMessage from dify_plugin import Tool +_NETWORK_ERROR = ( + "An error occurred while invoking the tool (network or upstream failure). " + "Please retry shortly." +) + class PoiSearchTool(Tool): def _invoke( @@ -18,28 +23,43 @@ def _invoke( if not keyword: yield self.create_text_message("Invalid parameter keyword") return - tk = self.runtime.credentials["tianditu_api_key"] - keyword_coords = requests.get( - geocoder_base_url - + "?ds=" - + json.dumps({"keyWord": keyword}, ensure_ascii=False) - + "&tk=" - + tk - ).json() - coords = ( - str(keyword_coords["location"]["lon"]) - + "," - + str(keyword_coords["location"]["lat"]) - ) - result = requests.get( - base_url - + "?center=" - + coords - + "&markers=" - + coords - + "&width=400&height=300&zoom=14&tk=" - + tk - ).content + + tk = self.runtime.credentials.get("tianditu_api_key") + if not tk: + yield self.create_text_message("Tianditu API key is required.") + return + + try: + geo_response = requests.get( + geocoder_base_url + + "?ds=" + + json.dumps({"keyWord": keyword}, ensure_ascii=False) + + "&tk=" + + tk, + timeout=10, + ) + geo_response.raise_for_status() + keyword_coords = geo_response.json() + coords = ( + str(keyword_coords["location"]["lon"]) + + "," + + str(keyword_coords["location"]["lat"]) + ) + img_response = requests.get( + base_url + + "?center=" + + coords + + "&markers=" + + coords + + "&width=400&height=300&zoom=14&tk=" + + tk, + timeout=10, + ) + img_response.raise_for_status() + result = img_response.content + except requests.exceptions.RequestException: + yield self.create_text_message(_NETWORK_ERROR) + return yield self.create_blob_message( blob=result, meta={"mime_type": "image/png"},