mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 18:15:02 +00:00
Test connection-backed tools, MCP refresh and API failure paths
Covers creating and editing tools whose key lives on a connection, the MCP action refresh, and the auth, ownership and error branches of the connections API. Editing a tool's key now also recovers a connection whose stored credentials can no longer be decrypted.
This commit is contained in:
1 parent
fff3e7c904
commit
4d533d24d6
4 files changed
+516
-2
No files matched your search
@@ -22,7 +22,7 @@ from docsgpt.api.user.team_sharing import effective_write_owner, visible_with_ac
|
||||
from docsgpt.connectors.catalog import definition_for_tool
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.core.url_validation import SSRFError, validate_url
|
||||
from docsgpt.security.encryption import decrypt_credentials, encrypt_credentials
|
||||
from docsgpt.security.encryption import CredentialDecryptionError, decrypt_credentials, encrypt_credentials
|
||||
from docsgpt.storage.db.base_repository import looks_like_uuid
|
||||
from docsgpt.storage.db.repositories.artifacts import ArtifactsRepository
|
||||
from docsgpt.storage.db.repositories.notes import NotesRepository
|
||||
@@ -491,7 +491,11 @@ def _update_connection_secrets(conn, user, tool_doc, config, config_requirements
|
||||
connection = ConnectorSessionsRepository(conn).get_for_user(str(tool_doc["connection_id"]), user)
|
||||
if connection is None:
|
||||
return None
|
||||
stored = connection_service.read_secrets(connection)
|
||||
try:
|
||||
stored = connection_service.read_secrets(connection)
|
||||
except CredentialDecryptionError:
|
||||
# Unreadable after a lost key: the new secret replaces it.
|
||||
stored = {}
|
||||
credentials = {**(stored.get("credentials") or {}), **secrets}
|
||||
connection_service.write_secrets(
|
||||
conn, connection, {**stored, "credentials": credentials},
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
"""Auth, ownership and failure paths of the connections API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from docsgpt.connectors import service
|
||||
from tests.connectors.test_connections_api import _call, _connection, _db, app # noqa: F401 (fixture)
|
||||
|
||||
_CID = "00000000-0000-0000-0000-000000000001"
|
||||
|
||||
|
||||
def _resources():
|
||||
from docsgpt.api.connector import connections as c
|
||||
|
||||
return [
|
||||
(c.ConnectorCatalog, "get", ()),
|
||||
(c.ConnectionsList, "get", ()),
|
||||
(c.ConnectionsList, "post", ()),
|
||||
(c.ConnectionDetail, "get", (_CID,)),
|
||||
(c.ConnectionDetail, "delete", (_CID,)),
|
||||
(c.ConnectionDisconnect, "post", (_CID,)),
|
||||
(c.ConnectionSetup, "post", (_CID,)),
|
||||
(c.ConnectionReconnect, "post", (_CID,)),
|
||||
(c.ConnectionPickerToken, "post", (_CID,)),
|
||||
(c.ConnectionClaim, "post", ()),
|
||||
(c.ConnectionToolPermissions, "put", (_CID, "t1")),
|
||||
(c.ConnectionRefreshTools, "post", (_CID,)),
|
||||
(c.ToolCredentialMode, "put", ("t1",)),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("index", range(13))
|
||||
def test_every_endpoint_needs_a_signed_in_user(app, index): # noqa: F811
|
||||
resource, method, args = _resources()[index]
|
||||
resp = _call(app, resource, method, "/api/connections", user=None, body={}, args=args)
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
||||
@pytest.mark.parametrize("index", [3, 4, 5, 6, 7, 8, 11])
|
||||
def test_someone_elses_connection_is_not_found(app, pg_conn, index): # noqa: F811
|
||||
resource, method, _ = _resources()[index]
|
||||
cid = _connection(pg_conn, user="victim", secrets={"credentials": {"token": "t"}})
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, resource, method, f"/api/connections/{cid}", body={"credentials": {"token": "x"}},
|
||||
args=(cid,))
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
class TestCreateErrors:
|
||||
def test_oauth_connector_does_not_take_pasted_keys(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionsList
|
||||
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, ConnectionsList, "post", "/api/connections",
|
||||
body={"connector_key": "google_drive", "credentials": {"token": "x"}})
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_credentials_must_be_an_object(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionsList
|
||||
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, ConnectionsList, "post", "/api/connections",
|
||||
body={"connector_key": "telegram", "credentials": "123:abc"})
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_disabled_connector_is_forbidden(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionsList
|
||||
from docsgpt.storage.db.repositories.connector_policies import ConnectorPoliciesRepository
|
||||
|
||||
ConnectorPoliciesRepository(pg_conn).upsert("telegram", enabled=False)
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, ConnectionsList, "post", "/api/connections",
|
||||
body={"connector_key": "telegram", "credentials": {"token": "123456:ABCDEFG"}})
|
||||
assert resp.status_code == 403
|
||||
assert resp.get_json()["code"] == "disabled"
|
||||
|
||||
def test_missing_field_is_a_bad_request(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionsList
|
||||
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, ConnectionsList, "post", "/api/connections",
|
||||
body={"connector_key": "telegram", "credentials": {}})
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
class TestToolPermissions:
|
||||
def _tool(self, pg_conn, cid):
|
||||
from sqlalchemy import text
|
||||
|
||||
return str(pg_conn.execute(text(
|
||||
"INSERT INTO user_tools (user_id, name, config, connection_id, actions) VALUES ('alice', 'telegram', "
|
||||
"'{}'::jsonb, CAST(:c AS uuid), '[{\"name\": \"telegram_send_message\", \"active\": true}]'::jsonb) "
|
||||
"RETURNING id"
|
||||
), {"c": cid}).scalar())
|
||||
|
||||
def test_sets_permissions_on_the_connection_tool(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionToolPermissions
|
||||
|
||||
cid = _connection(pg_conn, secrets={"credentials": {"token": "t"}})
|
||||
tool_id = self._tool(pg_conn, cid)
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, ConnectionToolPermissions, "put", "/p",
|
||||
body={"permissions": {"telegram_send_message": "off"}}, args=(cid, tool_id))
|
||||
assert resp.status_code == 200
|
||||
action = resp.get_json()["tool"]["actions"][0]
|
||||
assert action["permission"] == "off"
|
||||
|
||||
def test_unknown_permission_is_rejected(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionToolPermissions
|
||||
|
||||
cid = _connection(pg_conn, secrets={"credentials": {"token": "t"}})
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, ConnectionToolPermissions, "put", "/p",
|
||||
body={"permissions": {"telegram_send_message": "sometimes"}}, args=(cid, "t1"))
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_tool_of_another_connection_is_not_found(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionToolPermissions
|
||||
|
||||
cid = _connection(pg_conn, secrets={"credentials": {"token": "t"}})
|
||||
other = _connection(pg_conn, account_label="…zzzz", secrets={"credentials": {"token": "u"}})
|
||||
tool_id = self._tool(pg_conn, other)
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, ConnectionToolPermissions, "put", "/p",
|
||||
body={"permissions": {"telegram_send_message": "off"}}, args=(cid, tool_id))
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
class TestRefreshTools:
|
||||
def test_returns_the_diff(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionRefreshTools
|
||||
|
||||
cid = _connection(pg_conn, provider="mcp:https://m.example.com", auth_kind="mcp_oauth", secrets={})
|
||||
with _db(pg_conn), patch("docsgpt.connectors.mcp.refresh_mcp_tools",
|
||||
return_value={"added": ["a"], "removed": [], "tools": []}):
|
||||
resp = _call(app, ConnectionRefreshTools, "post", "/r", args=(cid,))
|
||||
assert resp.status_code == 200
|
||||
assert resp.get_json()["added"] == ["a"]
|
||||
|
||||
def test_signed_out_server_asks_to_reconnect(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionRefreshTools
|
||||
|
||||
cid = _connection(pg_conn, provider="mcp:https://m.example.com", auth_kind="mcp_oauth", secrets={})
|
||||
with _db(pg_conn), patch("docsgpt.connectors.mcp.refresh_mcp_tools",
|
||||
side_effect=service.ConnectionUnavailable("gone")):
|
||||
resp = _call(app, ConnectionRefreshTools, "post", "/r", args=(cid,))
|
||||
assert resp.status_code == 409
|
||||
assert resp.get_json()["code"] == "reconnect"
|
||||
|
||||
def test_server_failure_is_a_bad_gateway(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionRefreshTools
|
||||
|
||||
cid = _connection(pg_conn, provider="mcp:https://m.example.com", auth_kind="mcp_oauth", secrets={})
|
||||
with _db(pg_conn), patch("docsgpt.connectors.mcp.refresh_mcp_tools", side_effect=RuntimeError("boom")):
|
||||
resp = _call(app, ConnectionRefreshTools, "post", "/r", args=(cid,))
|
||||
assert resp.status_code == 502
|
||||
|
||||
|
||||
class TestPickerToken:
|
||||
def test_signed_out_connection_asks_to_reconnect(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionPickerToken
|
||||
|
||||
cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", secrets={})
|
||||
with _db(pg_conn), patch.object(service, "picker_token", side_effect=service.ConnectionUnavailable("x")):
|
||||
resp = _call(app, ConnectionPickerToken, "post", "/p", args=(cid,))
|
||||
assert resp.status_code == 409
|
||||
|
||||
def test_provider_outage_is_retryable(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionPickerToken
|
||||
|
||||
cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", secrets={})
|
||||
with _db(pg_conn), patch.object(service, "picker_token",
|
||||
side_effect=service.TransientConnectionError("x")):
|
||||
resp = _call(app, ConnectionPickerToken, "post", "/p", args=(cid,))
|
||||
assert resp.status_code == 503
|
||||
|
||||
def test_api_key_connection_has_no_picker(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionPickerToken
|
||||
|
||||
cid = _connection(pg_conn, secrets={"credentials": {"token": "t"}})
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, ConnectionPickerToken, "post", "/p", args=(cid,))
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
class TestClaim:
|
||||
def test_needs_provider_and_token(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionClaim
|
||||
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, ConnectionClaim, "post", "/c", body={"provider": "google_drive"})
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_unknown_token_is_not_found(self, app, pg_conn): # noqa: F811
|
||||
from docsgpt.api.connector.connections import ConnectionClaim
|
||||
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, ConnectionClaim, "post", "/c",
|
||||
body={"provider": "google_drive", "session_token": "nope"})
|
||||
assert resp.status_code == 404
|
||||
@@ -0,0 +1,122 @@
|
||||
"""Tests for re-reading an MCP connection's actions."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
import docsgpt.api.user # noqa: F401 (loads mcp_tool without the circular import)
|
||||
import docsgpt.agents.tools.mcp_tool # noqa: F401,E402 (patched below)
|
||||
from docsgpt.connectors import mcp, service
|
||||
from docsgpt.security.encryption import encrypt_json
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _db(conn):
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch.multiple("docsgpt.connectors.mcp", db_session=_yield, db_readonly=_yield), \
|
||||
patch.multiple("docsgpt.connectors.service", db_session=_yield, db_readonly=_yield):
|
||||
yield
|
||||
|
||||
|
||||
def _connection(conn, *, auth_kind="mcp_oauth", status="connected", secrets=None) -> dict:
|
||||
row = conn.execute(
|
||||
text(
|
||||
"INSERT INTO connector_sessions (user_id, provider, connector_key, auth_kind, status, account_label, "
|
||||
"server_url, encrypted_credentials) VALUES ('alice', 'mcp:https://mcp.linear.app', 'mcp:linear', :a, "
|
||||
":s, 'alice@example.com', 'https://mcp.linear.app', :e) RETURNING *"
|
||||
),
|
||||
{"a": auth_kind, "s": status, "e": encrypt_json(secrets or {}, "alice")},
|
||||
).one()
|
||||
return dict(row._mapping)
|
||||
|
||||
|
||||
def _tool(conn, connection_id, actions) -> str:
|
||||
return str(conn.execute(
|
||||
text(
|
||||
"INSERT INTO user_tools (user_id, name, config, actions, connection_id) VALUES ('alice', 'mcp_tool', "
|
||||
"CAST(:c AS jsonb), CAST(:a AS jsonb), CAST(:cid AS uuid)) RETURNING id"
|
||||
),
|
||||
{"c": json.dumps({"server_url": "https://mcp.linear.app/mcp", "auth_type": "oauth"}),
|
||||
"a": json.dumps(actions), "cid": str(connection_id)},
|
||||
).scalar())
|
||||
|
||||
|
||||
def _action(name, **extra):
|
||||
return {"name": name, "description": name, "parameters": {"properties": {"q": {"type": "string"}}}, **extra}
|
||||
|
||||
|
||||
class TestRefresh:
|
||||
def test_keeps_choices_adds_new_actions_and_drops_removed_ones(self, pg_conn):
|
||||
connection = _connection(pg_conn)
|
||||
tool_id = _tool(pg_conn, connection["id"], [
|
||||
_action("list_issues", active=False, require_approval=True),
|
||||
_action("archived_action"),
|
||||
])
|
||||
discovered = [
|
||||
_action("list_issues", annotations={"readOnlyHint": True}),
|
||||
_action("create_issue", annotations={"readOnlyHint": False}),
|
||||
]
|
||||
with _db(pg_conn), patch.object(mcp, "_discover", return_value=discovered):
|
||||
result = mcp.refresh_mcp_tools("alice", connection)
|
||||
assert result["added"] == ["create_issue"]
|
||||
assert result["removed"] == ["archived_action"]
|
||||
actions = {a["name"]: a for a in pg_conn.execute(
|
||||
text("SELECT actions FROM user_tools WHERE id = CAST(:i AS uuid)"), {"i": tool_id}
|
||||
).scalar()}
|
||||
assert set(actions) == {"list_issues", "create_issue"}
|
||||
# The user's choices survive a refresh.
|
||||
assert actions["list_issues"]["active"] is False
|
||||
assert actions["list_issues"]["require_approval"] is True
|
||||
# New writes default to needing approval.
|
||||
assert actions["create_issue"]["access"] == "write"
|
||||
assert actions["create_issue"]["require_approval"] is True
|
||||
assert actions["create_issue"]["parameters"]["properties"]["q"]["filled_by_llm"] is True
|
||||
assert [t["id"] for t in result["tools"]] == [tool_id]
|
||||
|
||||
def test_other_tools_on_the_connection_are_left_alone(self, pg_conn):
|
||||
connection = _connection(pg_conn)
|
||||
pg_conn.execute(
|
||||
text("INSERT INTO user_tools (user_id, name, config, connection_id) VALUES ('alice', 'telegram', "
|
||||
"'{}'::jsonb, CAST(:c AS uuid))"),
|
||||
{"c": str(connection["id"])},
|
||||
)
|
||||
with _db(pg_conn), patch.object(mcp, "_discover") as discover:
|
||||
result = mcp.refresh_mcp_tools("alice", connection)
|
||||
discover.assert_not_called()
|
||||
assert result == {"added": [], "removed": [], "tools": []}
|
||||
|
||||
|
||||
class TestDiscover:
|
||||
def test_oauth_connection_passes_its_id_not_tokens(self):
|
||||
connection = {"id": "c1", "auth_kind": "mcp_oauth", "status": "connected"}
|
||||
tool = {"config": {"server_url": "https://mcp.linear.app/mcp", "encrypted_credentials": "v1-blob"}}
|
||||
fake = MagicMock()
|
||||
fake.get_actions_metadata.return_value = [_action("search")]
|
||||
with patch("docsgpt.agents.tools.mcp_tool.MCPTool", return_value=fake) as tool_cls:
|
||||
assert mcp._discover("alice", connection, tool) == [_action("search")]
|
||||
config = tool_cls.call_args.args[0]
|
||||
assert config["connection_id"] == "c1"
|
||||
assert config["query_mode"] is True
|
||||
assert "encrypted_credentials" not in config
|
||||
fake.discover_tools.assert_called_once()
|
||||
|
||||
def test_api_key_connection_supplies_its_credentials(self):
|
||||
connection = {"id": "c1", "auth_kind": "api_key", "status": "connected"}
|
||||
with patch.object(service, "get_credentials", return_value={"bearer_token": "tok"}), \
|
||||
patch("docsgpt.agents.tools.mcp_tool.MCPTool") as tool_cls:
|
||||
mcp._discover("alice", connection, {"config": {"server_url": "https://m.example.com/mcp"}})
|
||||
assert tool_cls.call_args.args[0]["auth_credentials"] == {"bearer_token": "tok"}
|
||||
|
||||
def test_signed_out_oauth_connection_must_reconnect(self):
|
||||
connection = {"id": "c1", "auth_kind": "mcp_oauth", "status": "reconnect_needed"}
|
||||
with pytest.raises(service.ConnectionUnavailable), patch("docsgpt.agents.tools.mcp_tool.MCPTool") as tool_cls:
|
||||
mcp._discover("alice", connection, {"config": {}})
|
||||
tool_cls.assert_not_called()
|
||||
@@ -0,0 +1,185 @@
|
||||
"""Tests for creating and editing tools whose secret lives on a connection."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy import text
|
||||
|
||||
import docsgpt.api.user # noqa: F401 (import order: avoids the tools/tasks cycle)
|
||||
from docsgpt.connectors import service
|
||||
from docsgpt.security.encryption import encrypt_json
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app():
|
||||
return Flask(__name__)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _db(conn):
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch.multiple("docsgpt.api.user.tools.routes", db_session=_yield, db_readonly=_yield), \
|
||||
patch.multiple("docsgpt.connectors.service", db_session=_yield, db_readonly=_yield):
|
||||
yield
|
||||
|
||||
|
||||
def _call(app, resource, body, user="alice"):
|
||||
with app.test_request_context("/api/tools", method="POST", json=body):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": user}
|
||||
return resource().post()
|
||||
|
||||
|
||||
def _telegram(token="123456:SECRETTOKEN", **extra):
|
||||
return {
|
||||
"name": "telegram",
|
||||
"displayName": "Telegram",
|
||||
"description": "Send messages",
|
||||
"config": {"token": token},
|
||||
"status": True,
|
||||
**extra,
|
||||
}
|
||||
|
||||
|
||||
def _connection_row(conn, connection_id):
|
||||
return dict(conn.execute(
|
||||
text("SELECT * FROM connector_sessions WHERE id = CAST(:i AS uuid)"), {"i": connection_id}
|
||||
).one()._mapping)
|
||||
|
||||
|
||||
class TestCreateConnectedTool:
|
||||
def test_pasted_key_goes_on_a_connection_not_the_tool(self, app, pg_conn):
|
||||
from docsgpt.api.user.tools.routes import CreateTool
|
||||
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, CreateTool, _telegram())
|
||||
assert resp.status_code == 200
|
||||
body = resp.get_json()
|
||||
tool = dict(pg_conn.execute(
|
||||
text("SELECT * FROM user_tools WHERE id = CAST(:i AS uuid)"), {"i": body["id"]}
|
||||
).one()._mapping)
|
||||
assert str(tool["connection_id"]) == body["connection_id"]
|
||||
assert "SECRETTOKEN" not in str(tool["config"])
|
||||
assert "encrypted_credentials" not in (tool["config"] or {})
|
||||
row = _connection_row(pg_conn, body["connection_id"])
|
||||
assert service.read_secrets(row)["credentials"] == {"token": "123456:SECRETTOKEN"}
|
||||
|
||||
def test_same_key_twice_shares_one_connection(self, app, pg_conn):
|
||||
from docsgpt.api.user.tools.routes import CreateTool
|
||||
|
||||
with _db(pg_conn):
|
||||
first = _call(app, CreateTool, _telegram()).get_json()
|
||||
second = _call(app, CreateTool, _telegram()).get_json()
|
||||
assert first["id"] != second["id"]
|
||||
assert first["connection_id"] == second["connection_id"]
|
||||
|
||||
def test_existing_connection_is_used_by_id(self, app, pg_conn):
|
||||
from docsgpt.api.user.tools.routes import CreateTool
|
||||
|
||||
with _db(pg_conn):
|
||||
first = _call(app, CreateTool, _telegram()).get_json()
|
||||
resp = _call(app, CreateTool, _telegram(token="", connection_id=first["connection_id"]))
|
||||
assert resp.status_code == 200
|
||||
assert resp.get_json()["connection_id"] == first["connection_id"]
|
||||
|
||||
def test_someone_elses_connection_is_not_found(self, app, pg_conn):
|
||||
from docsgpt.api.user.tools.routes import CreateTool
|
||||
|
||||
with _db(pg_conn):
|
||||
victims = _call(app, CreateTool, _telegram(), user="victim").get_json()
|
||||
resp = _call(app, CreateTool, _telegram(token="", connection_id=victims["connection_id"]))
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_connection_of_another_service_is_not_found(self, app, pg_conn):
|
||||
from docsgpt.api.user.tools.routes import CreateTool
|
||||
|
||||
with _db(pg_conn):
|
||||
telegram = _call(app, CreateTool, _telegram()).get_json()
|
||||
resp = _call(app, CreateTool, {
|
||||
**_telegram(token=""), "name": "ntfy", "config": {}, "connection_id": telegram["connection_id"],
|
||||
})
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_missing_key_is_a_validation_error(self, app, pg_conn):
|
||||
from docsgpt.api.user.tools.routes import CreateTool
|
||||
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, CreateTool, _telegram(token=""))
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_disabled_connector_is_refused(self, app, pg_conn):
|
||||
from docsgpt.api.user.tools.routes import CreateTool
|
||||
from docsgpt.storage.db.repositories.connector_policies import ConnectorPoliciesRepository
|
||||
|
||||
ConnectorPoliciesRepository(pg_conn).upsert("telegram", enabled=False)
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, CreateTool, _telegram())
|
||||
assert resp.status_code == 403
|
||||
|
||||
def test_default_key_on_multi_user_falls_back_to_the_tool(self, app, pg_conn):
|
||||
from docsgpt.api.user.tools.routes import CreateTool
|
||||
|
||||
with _db(pg_conn), patch.object(service, "ensure_can_store_credentials",
|
||||
side_effect=service.EncryptionKeyNotConfigured("set a key")):
|
||||
resp = _call(app, CreateTool, _telegram())
|
||||
assert resp.status_code == 200
|
||||
tool = dict(pg_conn.execute(
|
||||
text("SELECT * FROM user_tools WHERE id = CAST(:i AS uuid)"), {"i": resp.get_json()["id"]}
|
||||
).one()._mapping)
|
||||
assert tool["connection_id"] is None
|
||||
assert "encrypted_credentials" in tool["config"]
|
||||
|
||||
|
||||
class TestUpdateConnectedTool:
|
||||
def _create(self, app, pg_conn, user="alice"):
|
||||
from docsgpt.api.user.tools.routes import CreateTool
|
||||
|
||||
with _db(pg_conn):
|
||||
return _call(app, CreateTool, _telegram(), user=user).get_json()
|
||||
|
||||
def test_new_key_is_written_to_the_connection(self, app, pg_conn):
|
||||
from docsgpt.api.user.tools.routes import UpdateTool
|
||||
|
||||
created = self._create(app, pg_conn)
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, UpdateTool, {"id": created["id"], "config": {"token": "999999:ROTATED"}})
|
||||
assert resp.status_code == 200
|
||||
row = _connection_row(pg_conn, created["connection_id"])
|
||||
assert service.read_secrets(row)["credentials"]["token"] == "999999:ROTATED"
|
||||
config = pg_conn.execute(
|
||||
text("SELECT config FROM user_tools WHERE id = CAST(:i AS uuid)"), {"i": created["id"]}
|
||||
).scalar()
|
||||
assert "ROTATED" not in str(config)
|
||||
|
||||
def test_new_key_replaces_unreadable_credentials(self, app, pg_conn):
|
||||
"""After a lost encryption key, editing the tool is a way back in."""
|
||||
from docsgpt.api.user.tools.routes import UpdateTool
|
||||
|
||||
created = self._create(app, pg_conn)
|
||||
pg_conn.execute(
|
||||
text("UPDATE connector_sessions SET encrypted_credentials = :b WHERE id = CAST(:i AS uuid)"),
|
||||
{"b": encrypt_json({"credentials": {"token": "x"}}, "someone-else"), "i": created["connection_id"]},
|
||||
)
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, UpdateTool, {"id": created["id"], "config": {"token": "999999:ROTATED"}})
|
||||
assert resp.status_code == 200
|
||||
row = _connection_row(pg_conn, created["connection_id"])
|
||||
assert service.read_secrets(row)["credentials"] == {"token": "999999:ROTATED"}
|
||||
|
||||
def test_editing_without_a_new_key_keeps_the_connection(self, app, pg_conn):
|
||||
from docsgpt.api.user.tools.routes import UpdateTool
|
||||
|
||||
created = self._create(app, pg_conn)
|
||||
with _db(pg_conn):
|
||||
resp = _call(app, UpdateTool, {"id": created["id"], "config": {}})
|
||||
assert resp.status_code == 200
|
||||
row = _connection_row(pg_conn, created["connection_id"])
|
||||
assert service.read_secrets(row)["credentials"]["token"] == "123456:SECRETTOKEN"
|
||||
Reference in new issue
Block a user