345 lines
13 KiB
Python
345 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import unittest
|
|
from collections.abc import Awaitable, Callable
|
|
from types import SimpleNamespace
|
|
from typing import Any, cast
|
|
|
|
from services.runtime.common import identity_assertion
|
|
from services.runtime.connector_ingestion import (
|
|
_invoke,
|
|
build_actor_resolver,
|
|
build_server,
|
|
load_hub_project_ids,
|
|
load_workspace_read_enabled,
|
|
)
|
|
from services.runtime.dispatch import service_module
|
|
from services.runtime.ingestion_auth import ActorContext
|
|
from services.runtime.ingestion_files import FileIngestionError
|
|
from services.runtime.ingestion_workflow import IngestionWorkflow
|
|
|
|
|
|
class _Workflow:
|
|
async def search_knowledge(self, **_: Any) -> dict[str, Any]:
|
|
return {}
|
|
|
|
async def get_knowledge_context(self, **_: Any) -> dict[str, Any]:
|
|
return {}
|
|
|
|
async def analyze_knowledge_files(self, **_: Any) -> dict[str, Any]:
|
|
return {}
|
|
|
|
async def submit_ingestion_proposal(self, **_: Any) -> dict[str, Any]:
|
|
return {}
|
|
|
|
async def get_ingestion_status(self, **_: Any) -> dict[str, Any]:
|
|
return {}
|
|
|
|
async def extract_knowledge_files(self, **_: Any) -> dict[str, Any]:
|
|
return {}
|
|
|
|
async def get_workspace_roots(self, **_: Any) -> dict[str, Any]:
|
|
return {}
|
|
|
|
async def get_workspace_children(self, **_: Any) -> dict[str, Any]:
|
|
return {}
|
|
|
|
async def get_workspace_page(self, **_: Any) -> dict[str, Any]:
|
|
return {}
|
|
|
|
async def create_text_source_proposal(self, **_: Any) -> dict[str, Any]:
|
|
return {}
|
|
|
|
|
|
async def _actor_resolver(_: Any) -> ActorContext:
|
|
return ActorContext(
|
|
subject="connector-client",
|
|
projects=frozenset({"11111111-1111-4111-8111-111111111111"}),
|
|
scopes=frozenset({"knowledge:read", "ingestion:write"}),
|
|
correlation_id="trace-test",
|
|
)
|
|
|
|
|
|
class IngestionConnectorTests(unittest.TestCase):
|
|
@staticmethod
|
|
def _resolve(
|
|
resolver: Callable[[Any], Awaitable[ActorContext]],
|
|
context: Any,
|
|
) -> ActorContext:
|
|
async def resolve() -> ActorContext:
|
|
return await resolver(context)
|
|
|
|
return asyncio.run(resolve())
|
|
|
|
def _context(
|
|
self,
|
|
*,
|
|
project: str = "knowledge-hub",
|
|
environment: str = "prod",
|
|
signing_key: bytes = b"k" * 32,
|
|
) -> Any:
|
|
client_id = "12345678-1234-4234-8234-123456789abc"
|
|
trace_id = "trace-dataplane-1"
|
|
headers = {
|
|
"authorization": "Bearer upstream-contextforge-token-is-not-the-caller-token",
|
|
"x-itworx-project": project,
|
|
"x-itworx-environment": environment,
|
|
"x-itworx-client-id": client_id,
|
|
"x-itworx-trace-id": trace_id,
|
|
}
|
|
headers["x-itworx-identity-assertion"] = identity_assertion(
|
|
signing_key,
|
|
project=project,
|
|
environment=environment,
|
|
client_id=client_id,
|
|
trace_id=trace_id,
|
|
)
|
|
context = SimpleNamespace(
|
|
request_context=SimpleNamespace(
|
|
request=SimpleNamespace(headers=headers),
|
|
)
|
|
)
|
|
return signing_key, context
|
|
|
|
def test_dedicated_dispatch_name_resolves_to_ingestion_runtime(self) -> None:
|
|
self.assertEqual(
|
|
service_module("connector-ingestion"),
|
|
"services.runtime.connector_ingestion",
|
|
)
|
|
|
|
def test_server_exposes_exactly_ten_ingestion_tools(self) -> None:
|
|
server = build_server(
|
|
workflow=cast(IngestionWorkflow, _Workflow()),
|
|
actor_resolver=_actor_resolver,
|
|
)
|
|
tools = asyncio.run(server.list_tools())
|
|
|
|
self.assertEqual(
|
|
{tool.name for tool in tools},
|
|
{
|
|
"search_knowledge",
|
|
"get_knowledge_context",
|
|
"analyze_knowledge_files",
|
|
"extract_knowledge_files",
|
|
"get_workspace_roots",
|
|
"get_workspace_children",
|
|
"get_workspace_page",
|
|
"create_text_source_proposal",
|
|
"submit_ingestion_proposal",
|
|
"get_ingestion_status",
|
|
},
|
|
)
|
|
|
|
def test_file_params_declares_chatgpt_files_array_contract(self) -> None:
|
|
server = build_server(
|
|
workflow=cast(IngestionWorkflow, _Workflow()),
|
|
actor_resolver=_actor_resolver,
|
|
)
|
|
tools = asyncio.run(server.list_tools())
|
|
analyze = next(
|
|
tool.model_dump(by_alias=True)
|
|
for tool in tools
|
|
if tool.name == "analyze_knowledge_files"
|
|
)
|
|
schema = analyze["inputSchema"]["properties"]["files"]
|
|
file_schema = schema["items"]
|
|
|
|
self.assertEqual(analyze["_meta"]["openai/fileParams"], ["files"])
|
|
self.assertEqual(
|
|
set(file_schema["properties"]),
|
|
{"download_url", "file_id", "mime_type", "file_name"},
|
|
)
|
|
self.assertEqual(set(file_schema["required"]), {"download_url", "file_id"})
|
|
self.assertEqual(file_schema["properties"]["download_url"]["type"], "string")
|
|
self.assertEqual(file_schema["properties"]["file_id"]["type"], "string")
|
|
self.assertEqual(file_schema["properties"]["mime_type"]["type"], "string")
|
|
self.assertEqual(file_schema["properties"]["file_name"]["type"], "string")
|
|
|
|
extract = next(
|
|
tool.model_dump(by_alias=True)
|
|
for tool in tools
|
|
if tool.name == "extract_knowledge_files"
|
|
)
|
|
self.assertEqual(extract["_meta"]["openai/fileParams"], ["files"])
|
|
|
|
def test_advertised_input_schemas_equal_governed_catalog(self) -> None:
|
|
from pathlib import Path
|
|
|
|
server = build_server(
|
|
workflow=cast(IngestionWorkflow, _Workflow()),
|
|
actor_resolver=_actor_resolver,
|
|
)
|
|
actual = {
|
|
tool.name: tool.model_dump(by_alias=True)["inputSchema"]
|
|
for tool in asyncio.run(server.list_tools())
|
|
}
|
|
catalog = Path(__file__).resolve().parents[1] / "catalog" / "tools"
|
|
for path in catalog.glob("ingestion-*.yaml"):
|
|
spec = json.loads(path.read_text(encoding="utf-8"))["spec"]
|
|
self.assertEqual(actual[spec["connectorToolName"]], spec["inputSchema"])
|
|
for name in (
|
|
"search_knowledge",
|
|
"extract_knowledge_files",
|
|
"analyze_knowledge_files",
|
|
"create_text_source_proposal",
|
|
"submit_ingestion_proposal",
|
|
"get_ingestion_status",
|
|
):
|
|
self.assertIn("project_id", actual[name]["properties"])
|
|
self.assertNotIn("project_id", actual[name].get("required", []))
|
|
self.assertIn("project_id", actual["get_knowledge_context"]["required"])
|
|
|
|
def test_write_tools_are_non_destructive_and_closed_world(self) -> None:
|
|
server = build_server(
|
|
workflow=cast(IngestionWorkflow, _Workflow()),
|
|
actor_resolver=_actor_resolver,
|
|
)
|
|
tools = {
|
|
tool.name: tool.model_dump(by_alias=True) for tool in asyncio.run(server.list_tools())
|
|
}
|
|
|
|
for name in ("analyze_knowledge_files", "submit_ingestion_proposal"):
|
|
annotations = tools[name]["annotations"]
|
|
self.assertFalse(annotations["readOnlyHint"])
|
|
self.assertFalse(annotations["destructiveHint"])
|
|
self.assertFalse(annotations["openWorldHint"])
|
|
self.assertTrue(tools["analyze_knowledge_files"]["annotations"]["idempotentHint"])
|
|
self.assertFalse(tools["submit_ingestion_proposal"]["annotations"]["idempotentHint"])
|
|
|
|
def test_file_error_audit_records_only_bounded_code_and_correlation(self) -> None:
|
|
async def reject_file(**_: Any) -> dict[str, Any]:
|
|
raise FileIngestionError(
|
|
"invalid_input", input_reason="attachment_file_id_invalid"
|
|
)
|
|
|
|
context = self._context()
|
|
private_file_values = [
|
|
{
|
|
"download_url": "private-download-reference",
|
|
"file_id": "private-file-reference",
|
|
}
|
|
]
|
|
with self.assertLogs("services.runtime.connector_ingestion", level="INFO") as captured:
|
|
result = asyncio.run(
|
|
_invoke(
|
|
"extract_knowledge_files",
|
|
context,
|
|
_actor_resolver,
|
|
reject_file,
|
|
files=private_file_values,
|
|
)
|
|
)
|
|
|
|
self.assertEqual(result["error"]["code"], "invalid_input")
|
|
audit_message = captured.records[0].getMessage().removeprefix("ingestion_tool ")
|
|
audit = json.loads(audit_message)
|
|
self.assertEqual(
|
|
set(audit),
|
|
{
|
|
"tool",
|
|
"result",
|
|
"status",
|
|
"latency_bucket",
|
|
"correlation_id",
|
|
"error_code",
|
|
"input_reason",
|
|
},
|
|
)
|
|
self.assertEqual(audit["error_code"], "invalid_input")
|
|
self.assertEqual(audit["input_reason"], "attachment_file_id_invalid")
|
|
self.assertEqual(audit["correlation_id"], "trace-test")
|
|
self.assertEqual(audit["tool"], "extract_knowledge_files")
|
|
self.assertEqual(audit["result"], "failure")
|
|
self.assertNotIn("private-download-reference", audit_message)
|
|
self.assertNotIn("private-file-reference", audit_message)
|
|
|
|
def test_dataplane_signed_assertion_is_trusted_and_authorization_is_ignored(self) -> None:
|
|
signing_key, context = self._context()
|
|
project_id = "11111111-1111-4111-8111-111111111111"
|
|
resolver = build_actor_resolver(
|
|
allowed_projects=frozenset({project_id}),
|
|
signing_key=signing_key,
|
|
)
|
|
|
|
actor = self._resolve(resolver, context)
|
|
|
|
self.assertEqual(actor.subject, "12345678-1234-4234-8234-123456789abc")
|
|
self.assertEqual(actor.projects, frozenset({project_id}))
|
|
self.assertEqual(actor.scopes, frozenset({"knowledge:read", "ingestion:write"}))
|
|
self.assertEqual(actor.correlation_id, "trace-dataplane-1")
|
|
self.assertEqual(actor.connector_project, "knowledge-hub")
|
|
self.assertEqual(actor.connector_environment, "prod")
|
|
|
|
def test_dataplane_assertion_rejects_invalid_signature_and_nonproduction(self) -> None:
|
|
signing_key, context = self._context()
|
|
headers = context.request_context.request.headers
|
|
headers["x-itworx-identity-assertion"] = "v1:0:" + "0" * 64
|
|
resolver = build_actor_resolver(
|
|
allowed_projects=frozenset({"11111111-1111-4111-8111-111111111111"}),
|
|
signing_key=signing_key,
|
|
)
|
|
with self.assertRaises(PermissionError):
|
|
self._resolve(resolver, context)
|
|
|
|
signing_key, staging = self._context(environment="staging")
|
|
resolver = build_actor_resolver(
|
|
allowed_projects=frozenset({"11111111-1111-4111-8111-111111111111"}),
|
|
signing_key=signing_key,
|
|
)
|
|
with self.assertRaises(PermissionError):
|
|
self._resolve(resolver, staging)
|
|
|
|
def test_shared_connector_key_cannot_forge_intake_identity_assertion(self) -> None:
|
|
shared_connector_key = b"s" * 32
|
|
_, shared_key_signed_context = self._context(signing_key=shared_connector_key)
|
|
dedicated_intake_key = b"i" * 32
|
|
resolver = build_actor_resolver(
|
|
allowed_projects=frozenset({"11111111-1111-4111-8111-111111111111"}),
|
|
signing_key=dedicated_intake_key,
|
|
)
|
|
with self.assertRaises(PermissionError):
|
|
self._resolve(resolver, shared_key_signed_context)
|
|
|
|
def test_target_project_allowlist_is_bounded_uuid_configuration(self) -> None:
|
|
project_id = "11111111-1111-4111-8111-111111111111"
|
|
self.assertEqual(
|
|
load_hub_project_ids(environ={"KNOWLEDGE_HUB_ALLOWED_PROJECT_IDS": f" {project_id} "}),
|
|
frozenset({project_id}),
|
|
)
|
|
for value in ("", "not-a-uuid", f"{project_id},{project_id}", ",".join([project_id] * 257)):
|
|
with self.subTest(value=value[:40]):
|
|
with self.assertRaises(ValueError):
|
|
load_hub_project_ids(environ={"KNOWLEDGE_HUB_ALLOWED_PROJECT_IDS": value})
|
|
|
|
def test_empty_allowlist_requires_workspace_opt_in_and_signed_exact_identity(self) -> None:
|
|
self.assertEqual(load_hub_project_ids(environ={}, workspace_enabled=True), frozenset())
|
|
signing_key, context = self._context()
|
|
resolver = build_actor_resolver(
|
|
allowed_projects=frozenset(), signing_key=signing_key, workspace_enabled=True
|
|
)
|
|
actor = self._resolve(resolver, context)
|
|
self.assertEqual(actor.projects, frozenset())
|
|
with self.assertRaises(ValueError):
|
|
build_actor_resolver(allowed_projects=frozenset(), signing_key=signing_key)
|
|
_, staging = self._context(environment="staging")
|
|
with self.assertRaises(PermissionError):
|
|
self._resolve(resolver, staging)
|
|
|
|
def test_workspace_read_configuration_is_default_off_and_strict(self) -> None:
|
|
self.assertFalse(load_workspace_read_enabled(environ={}))
|
|
self.assertTrue(
|
|
load_workspace_read_enabled(environ={"KNOWLEDGE_HUB_WORKSPACE_READ_ENABLED": "true"})
|
|
)
|
|
self.assertFalse(
|
|
load_workspace_read_enabled(environ={"KNOWLEDGE_HUB_WORKSPACE_READ_ENABLED": "false"})
|
|
)
|
|
for value in ("yes", "1", "true please"):
|
|
with self.subTest(value=value), self.assertRaises(ValueError):
|
|
load_workspace_read_enabled(environ={"KNOWLEDGE_HUB_WORKSPACE_READ_ENABLED": value})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|