693cd952cd
仓库首次入库(此前仅存在于本地工作区与测试服构建产物中,未上传)。 改动(src/knowledge_adapter/asset_content_processor.py): - 新增 _TRUNCATED_TRAILING_MARKDOWN_RESOURCE_IMAGE 规则,清除 chunk 末尾被截断的  - 回归测试 tests/test_asset_content_processor.py:: test_truncated_trailing_markdown_resource_image_is_removed 验证证据: - 测试服 120.76.217.51,镜像 knowledge-adapter-service:reviewed-test-aeab4b5fca2d (release aeab4b5fca2d8661):KB d015d05c-25c8-4d63-b2ab-e8590d3c2dfb 5 个查询全部 http=200、records=5、残留 resource:// = 0;KB b9dacd62 无退化;容器日志无 error/exception - 本地:pytest 141 项全通过、ruff 全通过 - 该文件与 258c31e0fe07fd08 源码的差异经逐文件 md5 比对,为唯一变更文件
633 lines
23 KiB
Python
633 lines
23 KiB
Python
import hashlib
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock
|
|
from urllib.parse import parse_qs, urlparse
|
|
|
|
import asyncpg
|
|
import httpx
|
|
import pytest
|
|
from alibabacloud_credentials.models import CredentialModel
|
|
from pydantic import SecretStr
|
|
|
|
from knowledge_adapter.factory import build_app, build_production_app
|
|
from knowledge_adapter.settings import Settings
|
|
|
|
DIFY_KEY = "f" * 32
|
|
KB_ID = "3904b40e-278b-4273-a6fa-db6cbc975587"
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def stub_weknora_api_key_file(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(
|
|
Settings,
|
|
"read_weknora_api_key",
|
|
lambda _settings: SecretStr("test-weknora-key"),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
async def configured_client() -> AsyncIterator[httpx.AsyncClient]:
|
|
settings = Settings(
|
|
DIFY_KEY_HASH_CURRENT=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
|
|
DIFY_KEY_HASH_NEXT="",
|
|
WEKNORA_BASE_URL="http://weknora.internal:8080",
|
|
WEKNORA_ALLOWED_KB_IDS=KB_ID,
|
|
WEKNORA_TENANT_ID=10008,
|
|
WEKNORA_OSS_STORAGE_BACKEND_ID="backend-test",
|
|
OSS_REGION="cn-shenzhen",
|
|
OSS_BUCKET="private-assets",
|
|
OSS_RAM_ROLE_NAME="knowledge-adapter-role",
|
|
ASSET_PUBLIC_BASE_URL="https://assets.example.com",
|
|
MAX_QUERY_CHARS=5,
|
|
)
|
|
|
|
def handler(_request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(200, json={"success": True, "data": []})
|
|
|
|
app = build_app(settings, weknora_transport=httpx.MockTransport(handler))
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
yield client
|
|
await app.state.weknora_client.aclose()
|
|
|
|
|
|
async def test_build_app_wires_settings_and_retrieval_service(
|
|
configured_client: httpx.AsyncClient,
|
|
) -> None:
|
|
response = await configured_client.post(
|
|
"/retrieval",
|
|
headers={"Authorization": f"Bearer {DIFY_KEY}"},
|
|
json={
|
|
"knowledge_id": KB_ID,
|
|
"query": "12345",
|
|
"retrieval_setting": {"top_k": 1, "score_threshold": 0},
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json() == {"records": []}
|
|
|
|
|
|
async def test_build_app_uses_configured_query_limit(
|
|
configured_client: httpx.AsyncClient,
|
|
) -> None:
|
|
response = await configured_client.post(
|
|
"/retrieval",
|
|
headers={"Authorization": f"Bearer {DIFY_KEY}"},
|
|
json={
|
|
"knowledge_id": KB_ID,
|
|
"query": "123456",
|
|
"retrieval_setting": {"top_k": 1, "score_threshold": 0},
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
|
|
|
|
async def test_app_lifespan_closes_weknora_client() -> None:
|
|
settings = Settings(
|
|
DIFY_KEY_HASH_CURRENT=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
|
|
DIFY_KEY_HASH_NEXT="",
|
|
WEKNORA_BASE_URL="http://weknora.internal:8080",
|
|
WEKNORA_ALLOWED_KB_IDS=KB_ID,
|
|
WEKNORA_TENANT_ID=10008,
|
|
WEKNORA_OSS_STORAGE_BACKEND_ID="backend-test",
|
|
OSS_REGION="cn-shenzhen",
|
|
OSS_BUCKET="private-assets",
|
|
OSS_RAM_ROLE_NAME="knowledge-adapter-role",
|
|
ASSET_PUBLIC_BASE_URL="https://assets.example.com",
|
|
)
|
|
app = build_app(
|
|
settings,
|
|
weknora_transport=httpx.MockTransport(
|
|
lambda _request: httpx.Response(200, json={"success": True, "data": []})
|
|
),
|
|
)
|
|
close_mock = AsyncMock()
|
|
app.state.weknora_client.aclose = close_mock
|
|
|
|
async with app.router.lifespan_context(app):
|
|
pass
|
|
|
|
close_mock.assert_awaited_once_with()
|
|
|
|
|
|
async def test_build_app_passes_configured_max_top_k_to_service() -> None:
|
|
settings = Settings(
|
|
DIFY_KEY_HASH_CURRENT=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
|
|
DIFY_KEY_HASH_NEXT="",
|
|
WEKNORA_BASE_URL="http://weknora.internal:8080",
|
|
WEKNORA_ALLOWED_KB_IDS=KB_ID,
|
|
WEKNORA_TENANT_ID=10008,
|
|
WEKNORA_OSS_STORAGE_BACKEND_ID="backend-test",
|
|
OSS_REGION="cn-shenzhen",
|
|
OSS_BUCKET="private-assets",
|
|
OSS_RAM_ROLE_NAME="knowledge-adapter-role",
|
|
ASSET_PUBLIC_BASE_URL="https://assets.example.com",
|
|
MAX_TOP_K=1,
|
|
)
|
|
base_chunk = {
|
|
"score": 0.9,
|
|
"knowledge_base_id": KB_ID,
|
|
"knowledge_id": "6f77a258-27ff-4278-a667-8fe242ff8e10",
|
|
"knowledge_title": "Manual",
|
|
"knowledge_filename": "manual.pdf",
|
|
"chunk_type": "text",
|
|
"image_info": "",
|
|
}
|
|
|
|
def handler(_request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"success": True,
|
|
"data": [
|
|
{
|
|
**base_chunk,
|
|
"id": "c7f61f89-19ce-45b1-ac60-dbeed5e2a491",
|
|
"content": "first",
|
|
"chunk_index": 0,
|
|
},
|
|
{
|
|
**base_chunk,
|
|
"id": "09ab721f-7d05-45de-ae0f-dac303276f6d",
|
|
"content": "second",
|
|
"chunk_index": 1,
|
|
},
|
|
],
|
|
},
|
|
)
|
|
|
|
app = build_app(settings, weknora_transport=httpx.MockTransport(handler))
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with app.router.lifespan_context(app):
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
response = await client.post(
|
|
"/retrieval",
|
|
headers={"Authorization": f"Bearer {DIFY_KEY}"},
|
|
json={
|
|
"knowledge_id": KB_ID,
|
|
"query": "query",
|
|
"retrieval_setting": {"top_k": 100, "score_threshold": 0},
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert [record["content"] for record in response.json()["records"]] == ["first"]
|
|
|
|
|
|
async def test_build_app_passes_configured_content_budgets_to_service() -> None:
|
|
settings = Settings(
|
|
DIFY_KEY_HASH_CURRENT=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
|
|
DIFY_KEY_HASH_NEXT="",
|
|
WEKNORA_BASE_URL="http://weknora.internal:8080",
|
|
WEKNORA_ALLOWED_KB_IDS=KB_ID,
|
|
WEKNORA_TENANT_ID=10008,
|
|
WEKNORA_OSS_STORAGE_BACKEND_ID="backend-test",
|
|
OSS_REGION="cn-shenzhen",
|
|
OSS_BUCKET="private-assets",
|
|
OSS_RAM_ROLE_NAME="knowledge-adapter-role",
|
|
ASSET_PUBLIC_BASE_URL="https://assets.example.com",
|
|
MAX_CHUNK_CHARS=50,
|
|
MAX_TOTAL_CHARS=70,
|
|
)
|
|
base_chunk = {
|
|
"score": 0.9,
|
|
"knowledge_base_id": KB_ID,
|
|
"knowledge_id": "6f77a258-27ff-4278-a667-8fe242ff8e10",
|
|
"knowledge_title": "Manual",
|
|
"knowledge_filename": "manual.pdf",
|
|
"chunk_type": "text",
|
|
"image_info": "",
|
|
}
|
|
|
|
def handler(_request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"success": True,
|
|
"data": [
|
|
{
|
|
**base_chunk,
|
|
"id": "c7f61f89-19ce-45b1-ac60-dbeed5e2a491",
|
|
"content": "a" * 100,
|
|
"chunk_index": 0,
|
|
},
|
|
{
|
|
**base_chunk,
|
|
"id": "09ab721f-7d05-45de-ae0f-dac303276f6d",
|
|
"content": "b" * 100,
|
|
"chunk_index": 1,
|
|
},
|
|
],
|
|
},
|
|
)
|
|
|
|
app = build_app(settings, weknora_transport=httpx.MockTransport(handler))
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with app.router.lifespan_context(app):
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
response = await client.post(
|
|
"/retrieval",
|
|
headers={"Authorization": f"Bearer {DIFY_KEY}"},
|
|
json={
|
|
"knowledge_id": KB_ID,
|
|
"query": "query",
|
|
"retrieval_setting": {"top_k": 5, "score_threshold": 0},
|
|
},
|
|
)
|
|
|
|
records = response.json()["records"]
|
|
assert response.status_code == 200
|
|
assert len(records) == 1
|
|
assert len(records[0]["content"]) == 50
|
|
|
|
|
|
class FakeRuntimePool:
|
|
def __init__(self, connection: object | None = None) -> None:
|
|
self.connection = connection
|
|
self.closed = False
|
|
|
|
@asynccontextmanager
|
|
async def acquire(self) -> AsyncIterator[object]:
|
|
if self.connection is None:
|
|
raise AssertionError("pool acquire was not expected")
|
|
yield self.connection
|
|
|
|
async def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class RecordingPoolFactory:
|
|
def __init__(self, pools: list[FakeRuntimePool] | None = None) -> None:
|
|
self.calls: list[dict[str, object]] = []
|
|
self.pools = pools or [FakeRuntimePool(), FakeRuntimePool()]
|
|
|
|
async def __call__(self, **kwargs: object) -> FakeRuntimePool:
|
|
self.calls.append(kwargs)
|
|
return self.pools[len(self.calls) - 1]
|
|
|
|
|
|
async def test_production_factory_creates_separate_bounded_pools(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
asset_secret = tmp_path / "asset-db-url"
|
|
weknora_secret = tmp_path / "weknora-db-url"
|
|
asset_secret.write_text("postgresql://adapter-test\n", encoding="utf-8")
|
|
weknora_secret.write_text("postgresql://weknora-readonly-test\n", encoding="utf-8")
|
|
settings = Settings(
|
|
DIFY_KEY_HASH_CURRENT=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
|
|
DIFY_KEY_HASH_NEXT="",
|
|
WEKNORA_BASE_URL="http://weknora.internal:8080",
|
|
WEKNORA_ALLOWED_KB_IDS=KB_ID,
|
|
WEKNORA_TENANT_ID=10008,
|
|
WEKNORA_OSS_STORAGE_BACKEND_ID="backend-test",
|
|
weknora_readonly_database_url_file=weknora_secret,
|
|
asset_database_url_file=asset_secret,
|
|
OSS_REGION="cn-shenzhen",
|
|
OSS_BUCKET="private-assets",
|
|
OSS_RAM_ROLE_NAME="knowledge-adapter-role",
|
|
ASSET_PUBLIC_BASE_URL="https://assets.example.com",
|
|
)
|
|
pool_factory = RecordingPoolFactory()
|
|
app = build_production_app(
|
|
settings,
|
|
pool_factory=pool_factory,
|
|
weknora_transport=httpx.MockTransport(
|
|
lambda _request: httpx.Response(200, json={"success": True, "data": []})
|
|
),
|
|
)
|
|
|
|
async with app.router.lifespan_context(app):
|
|
assert len(pool_factory.calls) == 2
|
|
asset_call, weknora_call = pool_factory.calls
|
|
assert asset_call == {
|
|
"dsn": "postgresql://adapter-test",
|
|
"min_size": 1,
|
|
"max_size": 5,
|
|
"command_timeout": 5,
|
|
"server_settings": {"statement_timeout": "5000"},
|
|
}
|
|
assert weknora_call == {
|
|
"dsn": "postgresql://weknora-readonly-test",
|
|
"min_size": 1,
|
|
"max_size": 2,
|
|
"command_timeout": 1,
|
|
"server_settings": {
|
|
"default_transaction_read_only": "on",
|
|
"statement_timeout": "1000",
|
|
},
|
|
}
|
|
assert not any(pool.closed for pool in pool_factory.pools)
|
|
|
|
assert all(pool.closed for pool in pool_factory.pools)
|
|
|
|
|
|
class RuntimeAdapterConnection:
|
|
def __init__(self) -> None:
|
|
self.queries: list[tuple[str, tuple[object, ...]]] = []
|
|
|
|
async def fetchrow(self, query: str, *args: object) -> dict[str, object] | None:
|
|
self.queries.append((query, args))
|
|
if query.lstrip().startswith("SELECT"):
|
|
return None
|
|
return {
|
|
"asset_id": "c7f61f89-19ce-45b1-ac60-dbeed5e2a491",
|
|
"source_ref": args[1],
|
|
"object_key": args[2],
|
|
"content_type": args[3],
|
|
"created_at": datetime(2026, 8, 24, tzinfo=UTC),
|
|
"status": "active",
|
|
}
|
|
|
|
|
|
class RuntimeWeKnoraConnection:
|
|
def __init__(self) -> None:
|
|
self.readonly_transactions: list[bool] = []
|
|
self.fetch_calls: list[tuple[str, tuple[object, ...]]] = []
|
|
|
|
@asynccontextmanager
|
|
async def transaction(self, *, readonly: bool) -> AsyncIterator[None]:
|
|
self.readonly_transactions.append(readonly)
|
|
yield
|
|
|
|
async def fetch(self, query: str, *args: object) -> list[dict[str, object]]:
|
|
self.fetch_calls.append((query, args))
|
|
return [
|
|
{
|
|
"tenant_id": 10008,
|
|
"storage_backend_id": "backend-test",
|
|
"provider": "oss",
|
|
"physical_path": (
|
|
"storage://backend-test/"
|
|
"oss://private-assets/weknora/10008/assets/image.png"
|
|
),
|
|
"mime_type": "image/png",
|
|
}
|
|
]
|
|
|
|
|
|
async def test_production_factory_wires_first_image_resolution_flow(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
source_ref = "resource://AbCdEfGhIjKlMnOpQrSt_1"
|
|
asset_secret = tmp_path / "asset-db-url"
|
|
weknora_secret = tmp_path / "weknora-db-url"
|
|
asset_secret.write_text("postgresql://adapter-test\n", encoding="utf-8")
|
|
weknora_secret.write_text("postgresql://weknora-readonly-test\n", encoding="utf-8")
|
|
settings = Settings(
|
|
DIFY_KEY_HASH_CURRENT=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
|
|
DIFY_KEY_HASH_NEXT="",
|
|
WEKNORA_BASE_URL="http://weknora.internal:8080",
|
|
WEKNORA_ALLOWED_KB_IDS=KB_ID,
|
|
WEKNORA_TENANT_ID=10008,
|
|
WEKNORA_OSS_STORAGE_BACKEND_ID="backend-test",
|
|
weknora_readonly_database_url_file=weknora_secret,
|
|
asset_database_url_file=asset_secret,
|
|
OSS_REGION="cn-shenzhen",
|
|
OSS_BUCKET="private-assets",
|
|
OSS_RAM_ROLE_NAME="knowledge-adapter-role",
|
|
ASSET_PUBLIC_BASE_URL="https://assets.example.com",
|
|
)
|
|
adapter_connection = RuntimeAdapterConnection()
|
|
weknora_connection = RuntimeWeKnoraConnection()
|
|
pool_factory = RecordingPoolFactory(
|
|
[FakeRuntimePool(adapter_connection), FakeRuntimePool(weknora_connection)]
|
|
)
|
|
|
|
def handler(_request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"success": True,
|
|
"data": [
|
|
{
|
|
"id": "09ab721f-7d05-45de-ae0f-dac303276f6d",
|
|
"content": f"",
|
|
"score": 0.9,
|
|
"knowledge_base_id": KB_ID,
|
|
"knowledge_id": "6f77a258-27ff-4278-a667-8fe242ff8e10",
|
|
"knowledge_title": "Manual",
|
|
"knowledge_filename": "manual.pdf",
|
|
"chunk_index": 0,
|
|
"chunk_type": "text",
|
|
"image_info": "[]",
|
|
}
|
|
],
|
|
},
|
|
)
|
|
|
|
app = build_production_app(
|
|
settings,
|
|
pool_factory=pool_factory,
|
|
weknora_transport=httpx.MockTransport(handler),
|
|
)
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with app.router.lifespan_context(app):
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
response = await client.post(
|
|
"/retrieval",
|
|
headers={"Authorization": f"Bearer {DIFY_KEY}"},
|
|
json={
|
|
"knowledge_id": KB_ID,
|
|
"query": "query",
|
|
"retrieval_setting": {"top_k": 1, "score_threshold": 0},
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["records"][0]["content"] == (
|
|
""
|
|
)
|
|
assert source_ref not in response.text
|
|
assert "weknora/10008/assets/image.png" not in response.text
|
|
assert len(adapter_connection.queries) == 2
|
|
assert weknora_connection.readonly_transactions == [True]
|
|
assert len(weknora_connection.fetch_calls) == 1
|
|
|
|
|
|
async def test_production_factory_defaults_to_asyncpg_pool_factory(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
asset_secret = tmp_path / "asset-db-url"
|
|
weknora_secret = tmp_path / "weknora-db-url"
|
|
asset_secret.write_text("postgresql://adapter-test\n", encoding="utf-8")
|
|
weknora_secret.write_text("postgresql://weknora-readonly-test\n", encoding="utf-8")
|
|
settings = Settings(
|
|
DIFY_KEY_HASH_CURRENT=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
|
|
DIFY_KEY_HASH_NEXT="",
|
|
WEKNORA_BASE_URL="http://weknora.internal:8080",
|
|
WEKNORA_ALLOWED_KB_IDS=KB_ID,
|
|
WEKNORA_TENANT_ID=10008,
|
|
WEKNORA_OSS_STORAGE_BACKEND_ID="backend-test",
|
|
weknora_readonly_database_url_file=weknora_secret,
|
|
asset_database_url_file=asset_secret,
|
|
OSS_REGION="cn-shenzhen",
|
|
OSS_BUCKET="private-assets",
|
|
OSS_RAM_ROLE_NAME="knowledge-adapter-role",
|
|
ASSET_PUBLIC_BASE_URL="https://assets.example.com",
|
|
)
|
|
pool_factory = RecordingPoolFactory()
|
|
monkeypatch.setattr(asyncpg, "create_pool", pool_factory)
|
|
|
|
app = build_production_app(
|
|
settings,
|
|
weknora_transport=httpx.MockTransport(
|
|
lambda _request: httpx.Response(200, json={"success": True, "data": []})
|
|
),
|
|
)
|
|
async with app.router.lifespan_context(app):
|
|
pass
|
|
|
|
assert len(pool_factory.calls) == 2
|
|
assert all(pool.closed for pool in pool_factory.pools)
|
|
|
|
|
|
class RuntimeAssetLookupConnection(RuntimeAdapterConnection):
|
|
async def fetchrow(self, query: str, *args: object) -> dict[str, object] | None:
|
|
self.queries.append((query, args))
|
|
return {
|
|
"asset_id": "c7f61f89-19ce-45b1-ac60-dbeed5e2a491",
|
|
"source_ref": "resource://AbCdEfGhIjKlMnOpQrSt_1",
|
|
"object_key": "weknora/10008/assets/image.png",
|
|
"content_type": "image/png",
|
|
"created_at": datetime(2026, 8, 24, tzinfo=UTC),
|
|
"status": "active",
|
|
}
|
|
|
|
|
|
class FactoryCredentialClient:
|
|
def get_credential(self) -> CredentialModel:
|
|
return CredentialModel(
|
|
access_key_id="test-access-key-id",
|
|
access_key_secret="test-access-key-secret",
|
|
security_token="test-security-token",
|
|
type="ecs_ram_role",
|
|
)
|
|
|
|
|
|
async def test_production_factory_wires_anonymous_asset_redirect(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
asset_secret = tmp_path / "asset-db-url"
|
|
weknora_secret = tmp_path / "weknora-db-url"
|
|
asset_secret.write_text("postgresql://adapter-test\n", encoding="utf-8")
|
|
weknora_secret.write_text("postgresql://weknora-readonly-test\n", encoding="utf-8")
|
|
settings = Settings(
|
|
DIFY_KEY_HASH_CURRENT=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
|
|
DIFY_KEY_HASH_NEXT="",
|
|
WEKNORA_BASE_URL="http://weknora.internal:8080",
|
|
WEKNORA_ALLOWED_KB_IDS=KB_ID,
|
|
WEKNORA_TENANT_ID=10008,
|
|
WEKNORA_OSS_STORAGE_BACKEND_ID="backend-test",
|
|
weknora_readonly_database_url_file=weknora_secret,
|
|
asset_database_url_file=asset_secret,
|
|
OSS_REGION="cn-shenzhen",
|
|
OSS_BUCKET="private-assets",
|
|
OSS_RAM_ROLE_NAME="knowledge-adapter-role",
|
|
ASSET_PUBLIC_BASE_URL="https://assets.example.com",
|
|
)
|
|
adapter_connection = RuntimeAssetLookupConnection()
|
|
pool_factory = RecordingPoolFactory(
|
|
[FakeRuntimePool(adapter_connection), FakeRuntimePool(RuntimeWeKnoraConnection())]
|
|
)
|
|
role_names: list[str] = []
|
|
|
|
def credential_client_factory(role_name: str) -> FactoryCredentialClient:
|
|
role_names.append(role_name)
|
|
return FactoryCredentialClient()
|
|
|
|
app = build_production_app(
|
|
settings,
|
|
pool_factory=pool_factory,
|
|
credential_client_factory=credential_client_factory,
|
|
weknora_transport=httpx.MockTransport(
|
|
lambda _request: httpx.Response(200, json={"success": True, "data": []})
|
|
),
|
|
)
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with app.router.lifespan_context(app):
|
|
async with httpx.AsyncClient(
|
|
transport=transport,
|
|
base_url="http://test",
|
|
follow_redirects=False,
|
|
) as client:
|
|
response = await client.get(
|
|
"/i/c7f61f89-19ce-45b1-ac60-dbeed5e2a491"
|
|
)
|
|
|
|
location = urlparse(response.headers["location"])
|
|
query = parse_qs(location.query)
|
|
assert response.status_code == 302
|
|
assert location.scheme == "https"
|
|
assert location.hostname == "private-assets.oss-cn-shenzhen.aliyuncs.com"
|
|
assert query["x-oss-signature-version"] == ["OSS4-HMAC-SHA256"]
|
|
assert query["x-oss-expires"] == ["600"]
|
|
assert role_names == ["knowledge-adapter-role"]
|
|
assert len(adapter_connection.queries) == 1
|
|
|
|
|
|
async def test_production_factory_uses_static_ak_secret_files(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
asset_secret = tmp_path / "asset-db-url"
|
|
weknora_secret = tmp_path / "weknora-db-url"
|
|
access_key_id_file = tmp_path / "oss-access-key-id"
|
|
access_key_secret_file = tmp_path / "oss-access-key-secret"
|
|
asset_secret.write_text("postgresql://adapter-test\n", encoding="utf-8")
|
|
weknora_secret.write_text("postgresql://weknora-readonly-test\n", encoding="utf-8")
|
|
access_key_id_file.write_text("test-access-key-id\n", encoding="utf-8")
|
|
access_key_secret_file.write_text("test-access-key-secret\n", encoding="utf-8")
|
|
settings = Settings(
|
|
DIFY_KEY_HASH_CURRENT=hashlib.sha256(DIFY_KEY.encode()).hexdigest(),
|
|
DIFY_KEY_HASH_NEXT="",
|
|
WEKNORA_BASE_URL="http://weknora.internal:8080",
|
|
WEKNORA_ALLOWED_KB_IDS=KB_ID,
|
|
WEKNORA_TENANT_ID=10008,
|
|
WEKNORA_OSS_STORAGE_BACKEND_ID="backend-test",
|
|
weknora_readonly_database_url_file=weknora_secret,
|
|
asset_database_url_file=asset_secret,
|
|
OSS_REGION="cn-shenzhen",
|
|
OSS_BUCKET="private-assets",
|
|
OSS_CREDENTIAL_MODE="static_ak",
|
|
OSS_ACCESS_KEY_ID_FILE=access_key_id_file,
|
|
OSS_ACCESS_KEY_SECRET_FILE=access_key_secret_file,
|
|
ASSET_PUBLIC_BASE_URL="https://assets.example.com",
|
|
)
|
|
adapter_connection = RuntimeAssetLookupConnection()
|
|
pool_factory = RecordingPoolFactory(
|
|
[FakeRuntimePool(adapter_connection), FakeRuntimePool(RuntimeWeKnoraConnection())]
|
|
)
|
|
|
|
def unexpected_ram_role_factory(_role_name: str) -> FactoryCredentialClient:
|
|
raise AssertionError("RAM role factory must not be used in static AK mode")
|
|
|
|
app = build_production_app(
|
|
settings,
|
|
pool_factory=pool_factory,
|
|
credential_client_factory=unexpected_ram_role_factory,
|
|
weknora_transport=httpx.MockTransport(
|
|
lambda _request: httpx.Response(200, json={"success": True, "data": []})
|
|
),
|
|
)
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with app.router.lifespan_context(app):
|
|
async with httpx.AsyncClient(
|
|
transport=transport,
|
|
base_url="http://test",
|
|
follow_redirects=False,
|
|
) as client:
|
|
response = await client.get("/i/c7f61f89-19ce-45b1-ac60-dbeed5e2a491")
|
|
|
|
location = urlparse(response.headers["location"])
|
|
query = parse_qs(location.query)
|
|
assert response.status_code == 302
|
|
assert location.hostname == "private-assets.oss-cn-shenzhen.aliyuncs.com"
|
|
assert query["x-oss-signature-version"] == ["OSS4-HMAC-SHA256"]
|
|
assert "x-oss-security-token" not in query
|