zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-23 08:56:42 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor: reformat all code for lefthook using ruff and gofmt (#16585)

Browse Source
This commit is contained in:
Wang Qi
2026-07-03 12:53:39 +08:00
committed by GitHub
parent 19fcb4a981
commit 6a4b9be426
588 changed files with 11123 additions and 15412 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
test/unit_test/api/apps/restful_apis/test_agentbots_access_control.py
View File

@@ -67,10 +67,11 @@ def _load_bot_api(monkeypatch, *, accessible, calls):
async def _gen():
yield 'data: {"event":"message","data":{"content":"ok"}}\n\n'
yield 'data: {"event":"message_end","data":{"content":"ok"}}\n\n'
return _gen()
_stub(monkeypatch, "quart", Response=lambda *a, **k: SimpleNamespace(headers=SimpleNamespace(add_header=lambda *aa, **kk: None)), request=SimpleNamespace())
_stub(monkeypatch, "api.apps", AUTH_BETA="beta", login_required=lambda *_a, **_k: (lambda func: func))
_stub(monkeypatch, "api.apps", AUTH_BETA="beta", login_required=lambda *_a, **_k: lambda func: func)
_stub(monkeypatch, "agent.canvas", Canvas=lambda *a, **k: SimpleNamespace(get_component_input_form=lambda _n: {}, get_prologue=lambda: "", get_mode=lambda: "agent"))
_stub(monkeypatch, "api.db.db_models", APIToken=SimpleNamespace(query=lambda **_k: [SimpleNamespace(tenant_id="attacker-tenant")]))
_stub(monkeypatch, "api.db.services.api_service", API4ConversationService=SimpleNamespace())

36
test/unit_test/api/apps/restful_apis/test_attachment_download_missing_blob.py
View File

@@ -64,7 +64,8 @@ def _load_agent_api(monkeypatch, *, storage_get):
return SimpleNamespace(payload=payload, headers={})
_stub(
monkeypatch, "api.apps",
monkeypatch,
"api.apps",
current_user=SimpleNamespace(id="tenant-1"),
login_required=lambda func: func,
)
@@ -72,7 +73,8 @@ def _load_agent_api(monkeypatch, *, storage_get):
_stub(monkeypatch, "api.db", CanvasCategory=SimpleNamespace())
_stub(monkeypatch, "api.db.db_models", Task=SimpleNamespace())
_stub(
monkeypatch, "api.db.services.api_service",
monkeypatch,
"api.db.services.api_service",
API4ConversationService=SimpleNamespace(
get_by_id=lambda _id: (False, None),
save=lambda **_k: True,
@@ -81,7 +83,8 @@ def _load_agent_api(monkeypatch, *, storage_get):
),
)
_stub(
monkeypatch, "api.db.services.canvas_service",
monkeypatch,
"api.db.services.canvas_service",
CanvasTemplateService=SimpleNamespace(),
UserCanvasService=SimpleNamespace(accessible=lambda *_a, **_k: True, query=lambda **_k: []),
completion=lambda *_a, **_k: None,
@@ -92,17 +95,23 @@ def _load_agent_api(monkeypatch, *, storage_get):
_stub(monkeypatch, "api.db.services.knowledgebase_service", KnowledgebaseService=SimpleNamespace())
_stub(monkeypatch, "api.db.services.pipeline_operation_log_service", PipelineOperationLogService=SimpleNamespace())
_stub(
monkeypatch, "api.db.services.task_service",
CANVAS_DEBUG_DOC_ID="", TaskService=SimpleNamespace(), queue_dataflow=lambda *_a, **_k: None,
monkeypatch,
"api.db.services.task_service",
CANVAS_DEBUG_DOC_ID="",
TaskService=SimpleNamespace(),
queue_dataflow=lambda *_a, **_k: None,
)
_stub(
monkeypatch, "api.db.services.user_service",
TenantService=SimpleNamespace(), UserService=SimpleNamespace(get_by_id=lambda *_a, **_k: (False, None)),
monkeypatch,
"api.db.services.user_service",
TenantService=SimpleNamespace(),
UserService=SimpleNamespace(get_by_id=lambda *_a, **_k: (False, None)),
)
_stub(monkeypatch, "api.db.services.user_canvas_version", UserCanvasVersionService=SimpleNamespace())
_stub(
monkeypatch, "api.utils.api_utils",
monkeypatch,
"api.utils.api_utils",
construct_json_result=lambda **kw: {"kind": "json", **kw},
get_data_error_result=lambda message="", code=0, data=False: {"kind": "data_error", "message": message},
get_error_data_result=lambda *_a, **_k: {"kind": "error"},
@@ -114,11 +123,13 @@ def _load_agent_api(monkeypatch, *, storage_get):
# Used as `@validate_request(...)` decorator factory at module level, so it
# must return an identity decorator (the lenient fallback would return None
# and `@None` raises TypeError during import).
validate_request=lambda *_a, **_k: (lambda func: func),
validate_request=lambda *_a, **_k: lambda func: func,
)
_stub(
monkeypatch, "common.settings",
retriever=SimpleNamespace(), kg_retriever=SimpleNamespace(),
monkeypatch,
"common.settings",
retriever=SimpleNamespace(),
kg_retriever=SimpleNamespace(),
# download_attachment reads settings.STORAGE_IMPL.get after
# `from common import settings` rebinds the module's `settings` name.
STORAGE_IMPL=SimpleNamespace(get=storage_get),
@@ -135,7 +146,8 @@ def _load_agent_api(monkeypatch, *, storage_get):
_stub(monkeypatch, "common.misc_utils", get_uuid=lambda: "uuid", thread_pool_exec=_thread_pool_exec)
_stub(
monkeypatch, "api.utils.web_utils",
monkeypatch,
"api.utils.web_utils",
CONTENT_TYPE_MAP={"markdown": "text/markdown"},
apply_safe_file_response_headers=lambda *_a, **_k: None,
)

9
test/unit_test/api/apps/restful_apis/test_openai_stream_no_duplicate.py
View File

@@ -80,7 +80,7 @@ def _load_openai_api(monkeypatch):
"api.utils.api_utils",
get_error_data_result=lambda *a, **k: {"code": 102},
get_request_json=lambda: {},
validate_request=lambda *_a, **_k: (lambda func: func),
validate_request=lambda *_a, **_k: lambda func: func,
)
_stub(monkeypatch, "common.constants", RetCode=SimpleNamespace(ARGUMENT_ERROR=102), StatusEnum=SimpleNamespace(VALID=SimpleNamespace(value="1")))
_stub(monkeypatch, "common.metadata_utils", convert_conditions=lambda *_a, **_k: None, meta_filter=lambda *_a, **_k: [])
@@ -107,6 +107,7 @@ async def _aiter(events):
def _collect_sse(module, events, **kwargs):
"""Run the SSE generator over `events` and return parsed JSON chunks
(the trailing `[DONE]` sentinel excluded)."""
async def run():
out = []
async for raw in module._stream_chat_completion_sse(_aiter(events), **kwargs):
@@ -205,11 +206,7 @@ def test_reasoning_content_streamed_separately(monkeypatch):
]
chunks = _collect_sse(module, events, need_reference=False, **_BASE_KWARGS)
reasoning = "".join(
c["choices"][0]["delta"].get("reasoning_content")
for c in chunks
if c != "[DONE]" and isinstance(c["choices"][0]["delta"].get("reasoning_content"), str)
)
reasoning = "".join(c["choices"][0]["delta"].get("reasoning_content") for c in chunks if c != "[DONE]" and isinstance(c["choices"][0]["delta"].get("reasoning_content"), str))
content = "".join(p for p in _content_pieces(chunks) if isinstance(p, str))
assert reasoning == "thinking"
assert content == "answer"

3
test/unit_test/api/apps/sdk/test_dify_retrieval.py
View File

@@ -1,4 +1,4 @@
#
#
# Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -84,6 +84,7 @@ class _FakeKGRetriever:
def _load_dify_retrieval(monkeypatch, *, kb, accessible, request_body, tenant_id, chunks=None):
"""Load dify_retrieval_api.py with minimum stubs to exercise the retrieval handler."""
def _add_tenant_id_to_kwargs(func):
async def wrapper(**kwargs):
kwargs["tenant_id"] = tenant_id

12
test/unit_test/api/apps/services/test_update_document_name_only.py
View File

@@ -99,16 +99,8 @@ def _load_update_document_name_only_module(monkeypatch, *, file_lookup):
fine_grained_tokenize=lambda tokens: tokens,
)
module_path = (
Path(__file__).resolve().parents[5]
/ "api"
/ "apps"
/ "services"
/ "document_api_service.py"
)
spec = importlib.util.spec_from_file_location(
"test_update_document_name_only_module", module_path
)
module_path = Path(__file__).resolve().parents[5] / "api" / "apps" / "services" / "document_api_service.py"
spec = importlib.util.spec_from_file_location("test_update_document_name_only_module", module_path)
module = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, "test_update_document_name_only_module", module)
spec.loader.exec_module(module)

8
test/unit_test/api/db/joint_services/test_tenant_model_service_max_tokens.py
View File

@@ -65,9 +65,7 @@ def test_max_tokens_falls_back_to_factory_when_model_extra_empty(monkeypatch):
],
)
config = tms.get_model_config_from_provider_instance(
"tenant-1", "chat", "gpt-test@default@OpenAI"
)
config = tms.get_model_config_from_provider_instance("tenant-1", "chat", "gpt-test@default@OpenAI")
assert config["max_tokens"] == 128000
@@ -115,8 +113,6 @@ def test_max_tokens_prefers_model_extra_over_factory(monkeypatch):
],
)
config = tms.get_model_config_from_provider_instance(
"tenant-1", "chat", "gpt-test@default@OpenAI"
)
config = tms.get_model_config_from_provider_instance("tenant-1", "chat", "gpt-test@default@OpenAI")
assert config["max_tokens"] == 32000

1
test/unit_test/api/db/services/test_dataset_access_permissions.py
View File

@@ -30,6 +30,7 @@ warnings.filterwarnings(
def _install_cv2_stub_if_unavailable():
try:
import cv2 # noqa: F401
return
except Exception:
pass

129
test/unit_test/api/db/services/test_dialog_service_final_answer.py
View File

@@ -48,6 +48,7 @@ warnings.filterwarnings(
def _install_cv2_stub_if_unavailable():
try:
import cv2 # noqa: F401
return
except Exception:
pass
@@ -197,6 +198,7 @@ class _FakeLangfuseClient:
def _collect(async_gen):
async def _run():
return [ev async for ev in async_gen]
return asyncio.run(_run())
@@ -204,6 +206,7 @@ def _collect(async_gen):
# Tests for async_ask (production code path)
# ---------------------------------------------------------------------------
@pytest.mark.p2
def test_async_ask_final_event_carries_decorated_answer(monkeypatch):
"""
@@ -218,23 +221,21 @@ def test_async_ask_final_event_carries_decorated_answer(monkeypatch):
chat_mdl = _StreamingChatModel(llm_answer)
retriever = _StubRetriever()
monkeypatch.setattr(dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB])
monkeypatch.setattr(
dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB]
)
monkeypatch.setattr(
dialog_service, "get_model_config_from_provider_instance",
dialog_service,
"get_model_config_from_provider_instance",
lambda _tid, _type, _name: _LLM_CONFIG,
)
monkeypatch.setattr(dialog_service, "LLMBundle", lambda _tid, _cfg: chat_mdl)
monkeypatch.setattr(dialog_service.settings, "retriever", retriever, raising=False)
monkeypatch.setattr(dialog_service.settings, "kg_retriever", retriever, raising=False)
monkeypatch.setattr(
dialog_service.DocMetadataService, "get_flatted_meta_by_kbs", lambda _ids: {}
)
monkeypatch.setattr(dialog_service.DocMetadataService, "get_flatted_meta_by_kbs", lambda _ids: {})
monkeypatch.setattr(dialog_service, "label_question", lambda _q, _kbs: "")
# kb_prompt calls DocumentService.get_by_ids which needs a live DB; stub it out.
monkeypatch.setattr(
dialog_service, "kb_prompt",
dialog_service,
"kb_prompt",
lambda _kbinfos, _max_tokens, **_kw: ["RAGFlow is a RAG engine."],
)
@@ -249,9 +250,7 @@ def test_async_ask_final_event_carries_decorated_answer(monkeypatch):
assert events, "async_ask must yield at least one event"
final_events = [e for e in events if e.get("final") is True]
assert len(final_events) == 1, (
f"Expected exactly one final event, got {len(final_events)}: {final_events}"
)
assert len(final_events) == 1, f"Expected exactly one final event, got {len(final_events)}: {final_events}"
final = final_events[0]
assert "answer" in final
@@ -267,22 +266,20 @@ def test_async_ask_delta_events_carry_incremental_text_only(monkeypatch):
chat_mdl = _StreamingChatModel("Incremental text for delta test.")
retriever = _StubRetriever()
monkeypatch.setattr(dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB])
monkeypatch.setattr(
dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB]
)
monkeypatch.setattr(
dialog_service, "get_model_config_from_provider_instance",
dialog_service,
"get_model_config_from_provider_instance",
lambda _tid, _type, _name: _LLM_CONFIG,
)
monkeypatch.setattr(dialog_service, "LLMBundle", lambda _tid, _cfg: chat_mdl)
monkeypatch.setattr(dialog_service.settings, "retriever", retriever, raising=False)
monkeypatch.setattr(dialog_service.settings, "kg_retriever", retriever, raising=False)
monkeypatch.setattr(
dialog_service.DocMetadataService, "get_flatted_meta_by_kbs", lambda _ids: {}
)
monkeypatch.setattr(dialog_service.DocMetadataService, "get_flatted_meta_by_kbs", lambda _ids: {})
monkeypatch.setattr(dialog_service, "label_question", lambda _q, _kbs: "")
monkeypatch.setattr(
dialog_service, "kb_prompt",
dialog_service,
"kb_prompt",
lambda _kbinfos, _max_tokens, **_kw: ["RAGFlow is a RAG engine."],
)
@@ -295,15 +292,13 @@ def test_async_ask_delta_events_carry_incremental_text_only(monkeypatch):
)
delta_events = [e for e in events if not e.get("final")]
final_events = [e for e in events if e.get("final") is True]
final_events = [e for e in events if e.get("final") is True]
assert len(final_events) == 1, f"Expected exactly one final event, got {len(final_events)}"
for ev in delta_events:
assert ev["reference"] == {}, f"Delta event must have empty reference, got: {ev['reference']}"
assert "chunks" in final_events[0]["reference"], (
"Final event reference must contain chunk data from decorate_answer()"
)
assert "chunks" in final_events[0]["reference"], "Final event reference must contain chunk data from decorate_answer()"
@pytest.mark.p2
@@ -311,9 +306,7 @@ def test_async_ask_empty_kb_ids_yields_error_final_event(monkeypatch):
"""
When kb_ids is empty, async_ask() must not crash with IndexError on kbs[0].
"""
monkeypatch.setattr(
dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: []
)
monkeypatch.setattr(dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [])
events = _collect(
dialog_service.async_ask(
@@ -357,6 +350,7 @@ def test_async_ask_stale_kb_ids_yields_error_final_event(monkeypatch):
# Tests for async_chat (production code path)
# ---------------------------------------------------------------------------
def _make_dialog(chat_mdl_stub):
"""Build a minimal dialog SimpleNamespace for async_chat()."""
return SimpleNamespace(
@@ -405,33 +399,30 @@ def test_async_chat_final_event_carries_decorated_answer(monkeypatch):
retriever = _StubRetriever()
# Stub out the heavy service/model calls
monkeypatch.setattr(dialog_service, "get_model_type_by_name", lambda _tid, _llm_id: ["chat"])
monkeypatch.setattr(
dialog_service, "get_model_type_by_name",
lambda _tid, _llm_id: ["chat"]
)
monkeypatch.setattr(
dialog_service, "get_model_config_from_provider_instance",
dialog_service,
"get_model_config_from_provider_instance",
lambda _tid, _type, _llm_id: _LLM_CONFIG,
)
monkeypatch.setattr(
dialog_service.TenantLangfuseService, "filter_by_tenant",
dialog_service.TenantLangfuseService,
"filter_by_tenant",
lambda tenant_id: None,
)
# get_models returns (kbs, embd_mdl, rerank_mdl, chat_mdl, tts_mdl)
monkeypatch.setattr(
dialog_service, "get_models",
dialog_service,
"get_models",
lambda _dialog, **_kwargs: ([_KB], chat_mdl, None, chat_mdl, None),
)
monkeypatch.setattr(
dialog_service.KnowledgebaseService, "get_field_map", lambda _kb_ids: {}
)
monkeypatch.setattr(
dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB]
)
monkeypatch.setattr(dialog_service.KnowledgebaseService, "get_field_map", lambda _kb_ids: {})
monkeypatch.setattr(dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB])
monkeypatch.setattr(dialog_service.settings, "retriever", retriever, raising=False)
monkeypatch.setattr(dialog_service, "label_question", lambda _q, _kbs: "")
monkeypatch.setattr(
dialog_service, "kb_prompt",
dialog_service,
"kb_prompt",
lambda _kbinfos, _max_tokens, **_kw: ["RAGFlow is a RAG engine."],
)
@@ -441,9 +432,7 @@ def test_async_chat_final_event_carries_decorated_answer(monkeypatch):
events = _collect(dialog_service.async_chat(dialog, messages, stream=True, quote=True))
final_events = [e for e in events if e.get("final") is True]
assert len(final_events) == 1, (
f"Expected exactly one final event, got {len(final_events)}: {final_events}"
)
assert len(final_events) == 1, f"Expected exactly one final event, got {len(final_events)}: {final_events}"
final = final_events[0]
assert "answer" in final
@@ -462,16 +451,15 @@ def test_async_chat_langfuse_uses_start_observation(monkeypatch):
chat_mdl = _StreamingChatModel(llm_answer)
retriever = _StubRetriever()
monkeypatch.setattr(dialog_service, "get_model_type_by_name", lambda _tid, _llm_id: ["chat"])
monkeypatch.setattr(
dialog_service, "get_model_type_by_name",
lambda _tid, _llm_id: ["chat"]
)
monkeypatch.setattr(
dialog_service, "get_model_config_from_provider_instance",
dialog_service,
"get_model_config_from_provider_instance",
lambda _tid, _type, _llm_id: _LLM_CONFIG,
)
monkeypatch.setattr(
dialog_service.TenantLangfuseService, "filter_by_tenant",
dialog_service.TenantLangfuseService,
"filter_by_tenant",
lambda tenant_id: SimpleNamespace(
public_key="public",
secret_key="secret",
@@ -486,12 +474,8 @@ def test_async_chat_langfuse_uses_start_observation(monkeypatch):
"get_models",
lambda _dialog, **_kwargs: ([_KB], chat_mdl, None, chat_mdl, None),
)
monkeypatch.setattr(
dialog_service.KnowledgebaseService, "get_field_map", lambda _kb_ids: {}
)
monkeypatch.setattr(
dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB]
)
monkeypatch.setattr(dialog_service.KnowledgebaseService, "get_field_map", lambda _kb_ids: {})
monkeypatch.setattr(dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB])
monkeypatch.setattr(dialog_service.settings, "retriever", retriever, raising=False)
monkeypatch.setattr(dialog_service, "label_question", lambda _q, _kbs: "")
monkeypatch.setattr(
@@ -530,17 +514,15 @@ def test_async_chat_langfuse_observation_includes_session_id(monkeypatch):
chat_mdl = _StreamingChatModel("Session traces should be grouped.")
retriever = _StubRetriever()
monkeypatch.setattr(
dialog_service, "get_model_type_by_name",
lambda _tid, _llm_id: ["chat"]
)
monkeypatch.setattr(dialog_service, "get_model_type_by_name", lambda _tid, _llm_id: ["chat"])
monkeypatch.setattr(
dialog_service,
"get_model_config_from_provider_instance",
lambda _tid, _type, _llm_id: _LLM_CONFIG,
)
monkeypatch.setattr(
dialog_service.TenantLangfuseService, "filter_by_tenant",
dialog_service.TenantLangfuseService,
"filter_by_tenant",
lambda tenant_id: SimpleNamespace(
public_key="public",
secret_key="secret",
@@ -554,12 +536,8 @@ def test_async_chat_langfuse_observation_includes_session_id(monkeypatch):
"get_models",
lambda _dialog, **_kwargs: ([_KB], chat_mdl, None, chat_mdl, None),
)
monkeypatch.setattr(
dialog_service.KnowledgebaseService, "get_field_map", lambda _kb_ids: {}
)
monkeypatch.setattr(
dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB]
)
monkeypatch.setattr(dialog_service.KnowledgebaseService, "get_field_map", lambda _kb_ids: {})
monkeypatch.setattr(dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB])
monkeypatch.setattr(dialog_service.settings, "retriever", retriever, raising=False)
monkeypatch.setattr(dialog_service, "label_question", lambda _q, _kbs: "")
monkeypatch.setattr(
@@ -635,16 +613,15 @@ def test_async_chat_continues_when_langfuse_observation_start_fails(monkeypatch)
chat_mdl = _StreamingChatModel(llm_answer)
retriever = _StubRetriever()
monkeypatch.setattr(dialog_service, "get_model_type_by_name", lambda _tid, _llm_id: ["chat"])
monkeypatch.setattr(
dialog_service, "get_model_type_by_name",
lambda _tid, _llm_id: ["chat"]
)
monkeypatch.setattr(
dialog_service, "get_model_config_from_provider_instance",
dialog_service,
"get_model_config_from_provider_instance",
lambda _tid, _type, _llm_id: _LLM_CONFIG,
)
monkeypatch.setattr(
dialog_service.TenantLangfuseService, "filter_by_tenant",
dialog_service.TenantLangfuseService,
"filter_by_tenant",
lambda tenant_id: SimpleNamespace(
public_key="public",
secret_key="secret",
@@ -659,12 +636,8 @@ def test_async_chat_continues_when_langfuse_observation_start_fails(monkeypatch)
"get_models",
lambda _dialog, **_kwargs: ([_KB], chat_mdl, None, chat_mdl, None),
)
monkeypatch.setattr(
dialog_service.KnowledgebaseService, "get_field_map", lambda _kb_ids: {}
)
monkeypatch.setattr(
dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB]
)
monkeypatch.setattr(dialog_service.KnowledgebaseService, "get_field_map", lambda _kb_ids: {})
monkeypatch.setattr(dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB])
monkeypatch.setattr(dialog_service.settings, "retriever", retriever, raising=False)
monkeypatch.setattr(dialog_service, "label_question", lambda _q, _kbs: "")
monkeypatch.setattr(

5
test/unit_test/api/db/services/test_dialog_service_use_sql_source_columns.py
View File

@@ -316,10 +316,7 @@ def test_async_chat_uses_all_docs_when_no_doc_ids_selected(monkeypatch):
)
monkeypatch.setattr(dialog_service.settings, "retriever", retriever, raising=False)
monkeypatch.setattr(
dialog_service, "get_model_type_by_name",
lambda _tid, _llm_id: ["chat"]
)
monkeypatch.setattr(dialog_service, "get_model_type_by_name", lambda _tid, _llm_id: ["chat"])
monkeypatch.setattr(
dialog_service,
"get_model_config_from_provider_instance",

25
test/unit_test/api/db/services/test_document_service_get_parsing_status.py
View File

@@ -31,6 +31,7 @@ warnings.filterwarnings(
def _install_cv2_stub_if_unavailable():
try:
import cv2 # noqa: F401
return
except Exception:
pass
@@ -70,6 +71,7 @@ from common.constants import TaskStatus # noqa: E402
# Helpers to access the original function bypassing @DB.connection_context()
# ---------------------------------------------------------------------------
def _unwrapped_get_parsing_status():
"""Return the original (un-decorated) get_parsing_status_by_kb_ids function.
@@ -84,6 +86,7 @@ def _unwrapped_get_parsing_status():
# Fake ORM helpers – mimic the minimal peewee query chain used by the function
# ---------------------------------------------------------------------------
class _FieldStub:
"""Minimal stand-in for a peewee model field used in select/where/group_by."""
@@ -130,6 +133,7 @@ def _make_fake_model(rows):
# Pytest fixture – patch DocumentService.model per test
# ---------------------------------------------------------------------------
@pytest.fixture()
def call_with_rows(monkeypatch):
"""Return a helper that runs get_parsing_status_by_kb_ids with fake DB rows."""
@@ -146,14 +150,11 @@ def call_with_rows(monkeypatch):
# Tests
# ---------------------------------------------------------------------------
_ALL_STATUS_FIELDS = frozenset(
["unstart_count", "running_count", "cancel_count", "done_count", "fail_count"]
)
_ALL_STATUS_FIELDS = frozenset(["unstart_count", "running_count", "cancel_count", "done_count", "fail_count"])
@pytest.mark.p2
class TestGetParsingStatusByKbIds:
# ------------------------------------------------------------------
# Edge-case: empty input list – must short-circuit before any DB call
# ------------------------------------------------------------------
@@ -224,16 +225,13 @@ class TestGetParsingStatusByKbIds:
def test_unknown_run_value_ignored(self, call_with_rows):
rows = [
{"kb_id": "kb-1", "run": "9", "cnt": 99}, # "9" is not a TaskStatus
{"kb_id": "kb-1", "run": "9", "cnt": 99}, # "9" is not a TaskStatus
{"kb_id": "kb-1", "run": TaskStatus.DONE.value, "cnt": 4},
]
result = call_with_rows(rows=rows, kb_ids=["kb-1"])
assert result["kb-1"]["done_count"] == 4
assert all(
result["kb-1"][f] == 0
for f in _ALL_STATUS_FIELDS - {"done_count"}
)
assert all(result["kb-1"][f] == 0 for f in _ALL_STATUS_FIELDS - {"done_count"})
# ------------------------------------------------------------------
# A row whose kb_id was NOT requested must not appear in the output
@@ -298,9 +296,7 @@ class TestGetParsingStatusByKbIds:
rows = [
{"kb_id": "kb-with-data", "run": TaskStatus.DONE.value, "cnt": 1},
]
result = call_with_rows(
rows=rows, kb_ids=["kb-with-data", "kb-empty-1", "kb-empty-2"]
)
result = call_with_rows(rows=rows, kb_ids=["kb-with-data", "kb-empty-1", "kb-empty-2"])
assert set(result.keys()) == {"kb-with-data", "kb-empty-1", "kb-empty-2"}
assert result["kb-empty-1"] == {f: 0 for f in _ALL_STATUS_FIELDS}
@@ -320,7 +316,4 @@ class TestGetParsingStatusByKbIds:
assert result["kb-1"]["done_count"] == 2
# SCHEDULE is not a tracked bucket
assert "schedule_count" not in result["kb-1"]
assert all(
result["kb-1"][f] == 0
for f in _ALL_STATUS_FIELDS - {"done_count"}
)
assert all(result["kb-1"][f] == 0 for f in _ALL_STATUS_FIELDS - {"done_count"})

7
test/unit_test/api/db/services/test_file_service_upload_document.py
View File

@@ -127,6 +127,7 @@ def test_upload_document_skips_cross_kb_document_id_collision(monkeypatch):
# Helpers shared by TestValidateUrlForCrawl
# ---------------------------------------------------------------------------
def _addrinfo(ip_str: str) -> list:
"""Build a minimal getaddrinfo-style result for a single address string."""
family = socket.AF_INET6 if ":" in ip_str else socket.AF_INET
@@ -137,6 +138,7 @@ def _addrinfo(ip_str: str) -> list:
# _validate_url_for_crawl SSRF-guard tests
# ---------------------------------------------------------------------------
@pytest.mark.p2
class TestValidateUrlForCrawl:
"""Focused regression suite for the SSRF guard on the URL-crawl path.
@@ -268,10 +270,7 @@ class TestValidateUrlForCrawl:
monkeypatch.setattr(
socket,
"getaddrinfo",
lambda h, p: (
_addrinfo("93.184.216.34")
+ _addrinfo("2606:2800:220:1:248:1893:25c8:1946")
),
lambda h, p: _addrinfo("93.184.216.34") + _addrinfo("2606:2800:220:1:248:1893:25c8:1946"),
)
hostname, resolved_ip = FileService._validate_url_for_crawl("https://example.com/")
assert hostname == "example.com"

45
test/unit_test/api/db/test_oceanbase_peewee.py
View File

@@ -20,22 +20,23 @@ class TestOceanBaseDatabase:
def test_oceanbase_in_pooled_database_enum(self):
"""Test that OCEANBASE is in PooledDatabase enum."""
assert hasattr(PooledDatabase, 'OCEANBASE')
assert hasattr(PooledDatabase, "OCEANBASE")
assert PooledDatabase.OCEANBASE.value == RetryingPooledOceanBaseDatabase
def test_oceanbase_in_database_lock_enum(self):
"""Test that OCEANBASE is in DatabaseLock enum."""
assert hasattr(DatabaseLock, 'OCEANBASE')
assert hasattr(DatabaseLock, "OCEANBASE")
def test_oceanbase_in_text_field_type_enum(self):
"""Test that OCEANBASE is in TextFieldType enum."""
assert hasattr(TextFieldType, 'OCEANBASE')
assert hasattr(TextFieldType, "OCEANBASE")
# OceanBase should use LONGTEXT like MySQL
assert TextFieldType.OCEANBASE.value == "LONGTEXT"
def test_oceanbase_database_inherits_mysql(self):
"""Test that OceanBase database inherits from PooledMySQLDatabase."""
from playhouse.pool import PooledMySQLDatabase
assert issubclass(RetryingPooledOceanBaseDatabase, PooledMySQLDatabase)
def test_oceanbase_database_init(self):
@@ -64,13 +65,13 @@ class TestOceanBaseDatabase:
def test_pooled_database_enum_values(self):
"""Test PooledDatabase enum has all expected values."""
expected = {'MYSQL', 'OCEANBASE', 'POSTGRES'}
expected = {"MYSQL", "OCEANBASE", "POSTGRES"}
actual = {e.name for e in PooledDatabase}
assert expected.issubset(actual), f"Missing: {expected - actual}"
def test_database_lock_enum_values(self):
"""Test DatabaseLock enum has all expected values."""
expected = {'MYSQL', 'OCEANBASE', 'POSTGRES'}
expected = {"MYSQL", "OCEANBASE", "POSTGRES"}
actual = set(DatabaseLock.__members__.keys())
assert expected.issubset(actual), f"Missing: {expected - actual}"
@@ -81,45 +82,49 @@ class TestOceanBaseConfiguration:
def test_settings_default_to_mysql(self):
"""Test that default DB_TYPE is mysql."""
import os
# Save original value
original = os.environ.get('DB_TYPE')
original = os.environ.get("DB_TYPE")
try:
# Remove DB_TYPE to test default
if 'DB_TYPE' in os.environ:
del os.environ['DB_TYPE']
if "DB_TYPE" in os.environ:
del os.environ["DB_TYPE"]
# Reload settings
from common import settings
settings.DATABASE_TYPE = os.getenv("DB_TYPE", "mysql")
assert settings.DATABASE_TYPE == "mysql"
finally:
# Restore original value
if original:
os.environ['DB_TYPE'] = original
os.environ["DB_TYPE"] = original
def test_settings_can_use_oceanbase(self):
"""Test that DB_TYPE can be set to oceanbase."""
import os
# Save original value
original = os.environ.get('DB_TYPE')
original = os.environ.get("DB_TYPE")
try:
os.environ['DB_TYPE'] = 'oceanbase'
os.environ["DB_TYPE"] = "oceanbase"
# Reload settings
from common import settings
settings.DATABASE_TYPE = os.getenv("DB_TYPE", "mysql")
assert settings.DATABASE_TYPE == "oceanbase"
finally:
# Restore original value
if original:
os.environ['DB_TYPE'] = original
os.environ["DB_TYPE"] = original
else:
if 'DB_TYPE' in os.environ:
del os.environ['DB_TYPE']
if "DB_TYPE" in os.environ:
del os.environ["DB_TYPE"]
if __name__ == "__main__":

35
test/unit_test/api/utils/test_doc_validation.py
View File

@@ -53,7 +53,7 @@ def test_validate_immutable_fields_no_changes():
doc.chunk_num = 10
doc.token_num = 100
doc.progress = 0.5
error_msg, error_code = validate_immutable_fields(update_doc_req, doc)
assert error_msg is None
assert error_code is None
@@ -66,7 +66,7 @@ def test_validate_immutable_fields_chunk_count_matches():
doc.chunk_num = 10
doc.token_num = 100
doc.progress = 0.5
error_msg, error_code = validate_immutable_fields(update_doc_req, doc)
assert error_msg is None
assert error_code is None
@@ -79,7 +79,7 @@ def test_validate_immutable_fields_token_count_matches():
doc.chunk_num = 10
doc.token_num = 100
doc.progress = 0.5
error_msg, error_code = validate_immutable_fields(update_doc_req, doc)
assert error_msg is None
assert error_code is None
@@ -92,7 +92,7 @@ def test_validate_immutable_fields_progress_matches():
doc.chunk_num = 10
doc.token_num = 100
doc.progress = 0.5
error_msg, error_code = validate_immutable_fields(update_doc_req, doc)
assert error_msg is None
assert error_code is None
@@ -105,7 +105,7 @@ def test_validate_immutable_fields_chunk_count_mismatch():
doc.chunk_num = 10
doc.token_num = 100
doc.progress = 0.5
error_msg, error_code = validate_immutable_fields(update_doc_req, doc)
assert error_msg == "Can't change `chunk_count`."
assert error_code == RetCode.DATA_ERROR
@@ -118,7 +118,7 @@ def test_validate_immutable_fields_token_count_mismatch():
doc.chunk_num = 10
doc.token_num = 100
doc.progress = 0.5
error_msg, error_code = validate_immutable_fields(update_doc_req, doc)
assert error_msg == "Can't change `token_count`."
assert error_code == RetCode.DATA_ERROR
@@ -131,7 +131,7 @@ def test_validate_immutable_fields_progress_mismatch():
doc.chunk_num = 10
doc.token_num = 100
doc.progress = 0.5
error_msg, error_code = validate_immutable_fields(update_doc_req, doc)
assert error_msg == "Can't change `progress`."
assert error_code == RetCode.DATA_ERROR
@@ -145,18 +145,18 @@ def test_validate_immutable_fields_progress_boundary_values():
doc.chunk_num = 10
doc.token_num = 100
doc.progress = 0.0
error_msg, error_code = validate_immutable_fields(update_doc_req, doc)
assert error_msg is None
assert error_code is None
# Test with 1.0
update_doc_req = UpdateDocumentReq(progress=1.0)
doc = Mock()
doc.chunk_num = 10
doc.token_num = 100
doc.progress = 1.0
error_msg, error_code = validate_immutable_fields(update_doc_req, doc)
assert error_msg is None
assert error_code is None
@@ -169,7 +169,7 @@ def test_validate_immutable_fields_none_values():
doc.chunk_num = 10
doc.token_num = 100
doc.progress = 0.5
error_msg, error_code = validate_immutable_fields(update_doc_req, doc)
assert error_msg is None
assert error_code is None
@@ -240,6 +240,7 @@ def test_validate_document_name_valid():
assert error_msg is None
assert error_code is None
def test_validate_document_name_attr_error():
"""Test valid document name update."""
req_doc_name = 0
@@ -312,7 +313,7 @@ def test_validate_chunk_method_valid():
doc = Mock()
doc.type = FileType.PDF
doc.name = "document.pdf"
error_msg, error_code = validate_chunk_method(doc)
assert error_msg is None
assert error_code is None
@@ -323,7 +324,7 @@ def test_validate_chunk_method_visual_not_supported():
doc = Mock()
doc.type = FileType.VISUAL
doc.name = "image.jpg"
error_msg, error_code = validate_chunk_method(doc)
assert "Not supported yet!" in error_msg
assert error_code == RetCode.DATA_ERROR
@@ -334,7 +335,7 @@ def test_validate_chunk_method_ppt_not_supported():
doc = Mock()
doc.type = FileType.PDF
doc.name = "presentation.ppt"
error_msg, error_code = validate_chunk_method(doc)
assert "Not supported yet!" in error_msg
assert error_code == RetCode.DATA_ERROR
@@ -345,7 +346,7 @@ def test_validate_chunk_method_pptx_not_supported():
doc = Mock()
doc.type = FileType.PDF
doc.name = "presentation.pptx"
error_msg, error_code = validate_chunk_method(doc)
assert "Not supported yet!" in error_msg
assert error_code == RetCode.DATA_ERROR
@@ -356,7 +357,7 @@ def test_validate_chunk_method_pages_not_supported():
doc = Mock()
doc.type = FileType.PDF
doc.name = "document.pages"
error_msg, error_code = validate_chunk_method(doc)
assert "Not supported yet!" in error_msg
assert error_code == RetCode.DATA_ERROR
@@ -367,7 +368,7 @@ def test_validate_chunk_method_other_extensions_still_valid():
doc = Mock()
doc.type = FileType.PDF
doc.name = "document.docx"
error_msg, error_code = validate_chunk_method(doc)
assert error_msg is None
assert error_code is None

13
test/unit_test/api/utils/test_health_utils_minio.py
View File

@@ -17,6 +17,7 @@
Unit tests for MinIO health check (check_minio_alive) and scheme/verify helpers.
Covers SSL/HTTPS and certificate verification (issues #13158, #13159).
"""
from unittest.mock import patch, Mock
@@ -27,6 +28,7 @@ class TestMinioSchemeAndVerify:
def test_scheme_http_when_secure_false(self, mock_settings):
mock_settings.MINIO = {"host": "minio:9000", "secure": False}
from api.utils.health_utils import _minio_scheme_and_verify
scheme, verify = _minio_scheme_and_verify()
assert scheme == "http"
assert verify is True
@@ -35,6 +37,7 @@ class TestMinioSchemeAndVerify:
def test_scheme_https_when_secure_true(self, mock_settings):
mock_settings.MINIO = {"host": "minio:9000", "secure": True}
from api.utils.health_utils import _minio_scheme_and_verify
scheme, verify = _minio_scheme_and_verify()
assert scheme == "https"
assert verify is True
@@ -43,6 +46,7 @@ class TestMinioSchemeAndVerify:
def test_scheme_https_when_secure_string_true(self, mock_settings):
mock_settings.MINIO = {"host": "minio:9000", "secure": "true"}
from api.utils.health_utils import _minio_scheme_and_verify
scheme, verify = _minio_scheme_and_verify()
assert scheme == "https"
@@ -50,6 +54,7 @@ class TestMinioSchemeAndVerify:
def test_verify_false_for_self_signed(self, mock_settings):
mock_settings.MINIO = {"host": "minio:9000", "secure": True, "verify": False}
from api.utils.health_utils import _minio_scheme_and_verify
scheme, verify = _minio_scheme_and_verify()
assert scheme == "https"
assert verify is False
@@ -58,6 +63,7 @@ class TestMinioSchemeAndVerify:
def test_verify_string_false(self, mock_settings):
mock_settings.MINIO = {"host": "minio:9000", "verify": "false"}
from api.utils.health_utils import _minio_scheme_and_verify
_, verify = _minio_scheme_and_verify()
assert verify is False
@@ -65,6 +71,7 @@ class TestMinioSchemeAndVerify:
def test_default_verify_true_when_key_missing(self, mock_settings):
mock_settings.MINIO = {"host": "minio:9000"}
from api.utils.health_utils import _minio_scheme_and_verify
_, verify = _minio_scheme_and_verify()
assert verify is True
@@ -80,6 +87,7 @@ class TestCheckMinioAlive:
mock_response.status_code = 200
mock_get.return_value = mock_response
from api.utils.health_utils import check_minio_alive
result = check_minio_alive()
assert result["status"] == "alive"
assert "elapsed" in result["message"]
@@ -96,6 +104,7 @@ class TestCheckMinioAlive:
mock_response.status_code = 200
mock_get.return_value = mock_response
from api.utils.health_utils import check_minio_alive
check_minio_alive()
call_args = mock_get.call_args
assert call_args[0][0] == "https://minio:9000/minio/health/live"
@@ -108,6 +117,7 @@ class TestCheckMinioAlive:
mock_response.status_code = 200
mock_get.return_value = mock_response
from api.utils.health_utils import check_minio_alive
check_minio_alive()
call_args = mock_get.call_args
assert call_args[1]["verify"] is False
@@ -120,6 +130,7 @@ class TestCheckMinioAlive:
mock_response.status_code = 503
mock_get.return_value = mock_response
from api.utils.health_utils import check_minio_alive
result = check_minio_alive()
assert result["status"] == "timeout"
@@ -129,6 +140,7 @@ class TestCheckMinioAlive:
mock_settings.MINIO = {"host": "minio:9000"}
mock_get.side_effect = ConnectionError("Connection refused")
from api.utils.health_utils import check_minio_alive
result = check_minio_alive()
assert result["status"] == "timeout"
assert "error" in result["message"]
@@ -141,6 +153,7 @@ class TestCheckMinioAlive:
mock_response.status_code = 200
mock_get.return_value = mock_response
from api.utils.health_utils import check_minio_alive
check_minio_alive()
call_args = mock_get.call_args
assert call_args[1]["timeout"] == 10

212
test/unit_test/api/utils/test_oceanbase_health.py
View File

@@ -16,6 +16,7 @@
"""
Unit tests for OceanBase health check and performance monitoring functionality.
"""
import inspect
import os
import types
@@ -27,20 +28,15 @@ from api.utils.health_utils import get_oceanbase_status, check_oceanbase_health
class TestOceanBaseHealthCheck:
"""Test cases for OceanBase health check functionality."""
@patch('api.utils.health_utils.OBConnection')
@patch.dict(os.environ, {'DOC_ENGINE': 'oceanbase'})
@patch("api.utils.health_utils.OBConnection")
@patch.dict(os.environ, {"DOC_ENGINE": "oceanbase"})
def test_get_oceanbase_status_success(self, mock_ob_class):
"""Test successful OceanBase status retrieval."""
# Setup mock
mock_ob_connection = Mock()
mock_ob_connection.uri = "localhost:2881"
mock_ob_connection.health.return_value = {
"uri": "localhost:2881",
"version_comment": "OceanBase 4.3.5.1",
"status": "healthy",
"connection": "connected"
}
mock_ob_connection.health.return_value = {"uri": "localhost:2881", "version_comment": "OceanBase 4.3.5.1", "status": "healthy", "connection": "connected"}
mock_ob_connection.get_performance_metrics.return_value = {
"connection": "connected",
"latency_ms": 5.2,
@@ -49,13 +45,13 @@ class TestOceanBaseHealthCheck:
"query_per_second": 150,
"slow_queries": 2,
"active_connections": 10,
"max_connections": 300
"max_connections": 300,
}
mock_ob_class.return_value = mock_ob_connection
# Execute
result = get_oceanbase_status()
# Assert
assert result["status"] == "alive"
assert "message" in result
@@ -63,36 +59,31 @@ class TestOceanBaseHealthCheck:
assert "performance" in result["message"]
assert result["message"]["health"]["status"] == "healthy"
assert result["message"]["performance"]["latency_ms"] == 5.2
@patch.dict(os.environ, {'DOC_ENGINE': 'elasticsearch'})
@patch.dict(os.environ, {"DOC_ENGINE": "elasticsearch"})
def test_get_oceanbase_status_not_configured(self):
"""Test OceanBase status when not configured."""
with pytest.raises(Exception) as exc_info:
get_oceanbase_status()
assert "OceanBase is not in use" in str(exc_info.value)
@patch('api.utils.health_utils.OBConnection')
@patch.dict(os.environ, {'DOC_ENGINE': 'oceanbase'})
@patch("api.utils.health_utils.OBConnection")
@patch.dict(os.environ, {"DOC_ENGINE": "oceanbase"})
def test_get_oceanbase_status_connection_error(self, mock_ob_class):
"""Test OceanBase status when connection fails."""
mock_ob_class.side_effect = Exception("Connection failed")
result = get_oceanbase_status()
assert result["status"] == "timeout"
assert "error" in result["message"]
@patch('api.utils.health_utils.OBConnection')
@patch.dict(os.environ, {'DOC_ENGINE': 'oceanbase'})
@patch("api.utils.health_utils.OBConnection")
@patch.dict(os.environ, {"DOC_ENGINE": "oceanbase"})
def test_check_oceanbase_health_healthy(self, mock_ob_class):
"""Test OceanBase health check returns healthy status."""
mock_ob_connection = Mock()
mock_ob_connection.health.return_value = {
"uri": "localhost:2881",
"version_comment": "OceanBase 4.3.5.1",
"status": "healthy",
"connection": "connected"
}
mock_ob_connection.health.return_value = {"uri": "localhost:2881", "version_comment": "OceanBase 4.3.5.1", "status": "healthy", "connection": "connected"}
mock_ob_connection.get_performance_metrics.return_value = {
"connection": "connected",
"latency_ms": 5.2,
@@ -101,28 +92,23 @@ class TestOceanBaseHealthCheck:
"query_per_second": 150,
"slow_queries": 0,
"active_connections": 10,
"max_connections": 300
"max_connections": 300,
}
mock_ob_class.return_value = mock_ob_connection
result = check_oceanbase_health()
assert result["status"] == "healthy"
assert result["details"]["connection"] == "connected"
assert result["details"]["latency_ms"] == 5.2
assert result["details"]["query_per_second"] == 150
@patch('api.utils.health_utils.OBConnection')
@patch.dict(os.environ, {'DOC_ENGINE': 'oceanbase'})
@patch("api.utils.health_utils.OBConnection")
@patch.dict(os.environ, {"DOC_ENGINE": "oceanbase"})
def test_check_oceanbase_health_degraded(self, mock_ob_class):
"""Test OceanBase health check returns degraded status for high latency."""
mock_ob_connection = Mock()
mock_ob_connection.health.return_value = {
"uri": "localhost:2881",
"version_comment": "OceanBase 4.3.5.1",
"status": "healthy",
"connection": "connected"
}
mock_ob_connection.health.return_value = {"uri": "localhost:2881", "version_comment": "OceanBase 4.3.5.1", "status": "healthy", "connection": "connected"}
mock_ob_connection.get_performance_metrics.return_value = {
"connection": "connected",
"latency_ms": 1500.0, # High latency > 1000ms
@@ -131,43 +117,35 @@ class TestOceanBaseHealthCheck:
"query_per_second": 50,
"slow_queries": 5,
"active_connections": 10,
"max_connections": 300
"max_connections": 300,
}
mock_ob_class.return_value = mock_ob_connection
result = check_oceanbase_health()
assert result["status"] == "degraded"
assert result["details"]["latency_ms"] == 1500.0
@patch('api.utils.health_utils.OBConnection')
@patch.dict(os.environ, {'DOC_ENGINE': 'oceanbase'})
@patch("api.utils.health_utils.OBConnection")
@patch.dict(os.environ, {"DOC_ENGINE": "oceanbase"})
def test_check_oceanbase_health_unhealthy(self, mock_ob_class):
"""Test OceanBase health check returns unhealthy status."""
mock_ob_connection = Mock()
mock_ob_connection.health.return_value = {
"uri": "localhost:2881",
"status": "unhealthy",
"connection": "disconnected",
"error": "Connection timeout"
}
mock_ob_connection.get_performance_metrics.return_value = {
"connection": "disconnected",
"error": "Connection timeout"
}
mock_ob_connection.health.return_value = {"uri": "localhost:2881", "status": "unhealthy", "connection": "disconnected", "error": "Connection timeout"}
mock_ob_connection.get_performance_metrics.return_value = {"connection": "disconnected", "error": "Connection timeout"}
mock_ob_class.return_value = mock_ob_connection
result = check_oceanbase_health()
assert result["status"] == "unhealthy"
assert result["details"]["connection"] == "disconnected"
assert "error" in result["details"]
@patch.dict(os.environ, {'DOC_ENGINE': 'elasticsearch'})
@patch.dict(os.environ, {"DOC_ENGINE": "elasticsearch"})
def test_check_oceanbase_health_not_configured(self):
"""Test OceanBase health check when not configured."""
result = check_oceanbase_health()
assert result["status"] == "not_configured"
assert result["details"]["connection"] == "not_configured"
assert "not configured" in result["details"]["message"].lower()
@@ -175,29 +153,32 @@ class TestOceanBaseHealthCheck:
class TestOBConnectionPerformanceMetrics:
"""Test cases for OBConnection performance metrics methods."""
def _create_mock_connection(self):
"""Create a mock OBConnection with actual methods."""
# Create a simple object and bind the real methods to it
class MockConn:
pass
conn = MockConn()
# Get the actual class from the singleton wrapper's closure
from rag.utils import ob_conn
# OBConnection is wrapped by @singleton decorator, so it's a function
# The original class is stored in the closure of the singleton function
# Find the class by checking all closure cells
ob_connection_class = None
if hasattr(ob_conn.OBConnection, '__closure__') and ob_conn.OBConnection.__closure__:
if hasattr(ob_conn.OBConnection, "__closure__") and ob_conn.OBConnection.__closure__:
for cell in ob_conn.OBConnection.__closure__:
cell_value = cell.cell_contents
if inspect.isclass(cell_value):
ob_connection_class = cell_value
break
if ob_connection_class is None:
raise ValueError("Could not find OBConnection class in closure")
# Bind the actual methods to our mock object
conn.get_performance_metrics = types.MethodType(ob_connection_class.get_performance_metrics, conn)
conn._get_storage_info = types.MethodType(ob_connection_class._get_storage_info, conn)
@@ -205,7 +186,7 @@ class TestOBConnectionPerformanceMetrics:
conn._get_slow_query_count = types.MethodType(ob_connection_class._get_slow_query_count, conn)
conn._estimate_qps = types.MethodType(ob_connection_class._estimate_qps, conn)
return conn
def test_get_performance_metrics_success(self):
"""Test successful retrieval of performance metrics."""
# Create mock connection with actual methods
@@ -214,29 +195,27 @@ class TestOBConnectionPerformanceMetrics:
conn.client = mock_client
conn.uri = "localhost:2881"
conn.db_name = "test"
# Mock client methods - create separate mock results for each call
mock_result1 = Mock()
mock_result1.fetchone.return_value = (1,)
mock_result2 = Mock()
mock_result2.fetchone.return_value = (100.5,)
mock_result3 = Mock()
mock_result3.fetchone.return_value = (100.0,)
mock_result4 = Mock()
mock_result4.fetchall.return_value = [
(1, 'user', 'host', 'db', 'Query', 0, 'executing', 'SELECT 1')
]
mock_result4.fetchone.return_value = ('max_connections', '300')
mock_result4.fetchall.return_value = [(1, "user", "host", "db", "Query", 0, "executing", "SELECT 1")]
mock_result4.fetchone.return_value = ("max_connections", "300")
mock_result5 = Mock()
mock_result5.fetchone.return_value = (0,)
mock_result6 = Mock()
mock_result6.fetchone.return_value = (5,)
# Setup side_effect to return different mocks for different queries
def sql_side_effect(query):
if "SELECT 1" in query:
@@ -254,21 +233,22 @@ class TestOBConnectionPerformanceMetrics:
elif "information_schema.processlist" in query and "COUNT" in query:
return mock_result6
return Mock()
mock_client.perform_raw_text_sql.side_effect = sql_side_effect
mock_client.pool_size = 300
# Mock logger
import logging
conn.logger = logging.getLogger('test')
conn.logger = logging.getLogger("test")
result = conn.get_performance_metrics()
assert result["connection"] == "connected"
assert result["latency_ms"] >= 0
assert "storage_used" in result
assert "storage_total" in result
def test_get_performance_metrics_connection_error(self):
"""Test performance metrics when connection fails."""
# Create mock connection with actual methods
@@ -277,14 +257,14 @@ class TestOBConnectionPerformanceMetrics:
conn.client = mock_client
conn.uri = "localhost:2881"
conn.logger = Mock()
mock_client.perform_raw_text_sql.side_effect = Exception("Connection failed")
result = conn.get_performance_metrics()
assert result["connection"] == "disconnected"
assert "error" in result
def test_get_storage_info_success(self):
"""Test successful retrieval of storage information."""
# Create mock connection with actual methods
@@ -293,27 +273,27 @@ class TestOBConnectionPerformanceMetrics:
conn.client = mock_client
conn.db_name = "test"
conn.logger = Mock()
mock_result1 = Mock()
mock_result1.fetchone.return_value = (100.5,)
mock_result2 = Mock()
mock_result2.fetchone.return_value = (100.0,)
def sql_side_effect(query):
if "information_schema.tables" in query:
return mock_result1
elif "__all_disk_stat" in query:
return mock_result2
return Mock()
mock_client.perform_raw_text_sql.side_effect = sql_side_effect
result = conn._get_storage_info()
assert "storage_used" in result
assert "storage_total" in result
assert "MB" in result["storage_used"]
def test_get_storage_info_fallback(self):
"""Test storage info with fallback when total space unavailable."""
# Create mock connection with actual methods
@@ -322,7 +302,7 @@ class TestOBConnectionPerformanceMetrics:
conn.client = mock_client
conn.db_name = "test"
conn.logger = Mock()
# First query succeeds, second fails
def side_effect(query):
if "information_schema.tables" in query:
@@ -331,14 +311,14 @@ class TestOBConnectionPerformanceMetrics:
return mock_result
else:
raise Exception("Table not found")
mock_client.perform_raw_text_sql.side_effect = side_effect
result = conn._get_storage_info()
assert "storage_used" in result
assert "storage_total" in result
def test_get_connection_pool_stats(self):
"""Test retrieval of connection pool statistics."""
# Create mock connection with actual methods
@@ -347,31 +327,28 @@ class TestOBConnectionPerformanceMetrics:
conn.client = mock_client
conn.logger = Mock()
mock_client.pool_size = 300
mock_result1 = Mock()
mock_result1.fetchall.return_value = [
(1, 'user', 'host', 'db', 'Query', 0, 'executing', 'SELECT 1'),
(2, 'user', 'host', 'db', 'Sleep', 10, None, None)
]
mock_result1.fetchall.return_value = [(1, "user", "host", "db", "Query", 0, "executing", "SELECT 1"), (2, "user", "host", "db", "Sleep", 10, None, None)]
mock_result2 = Mock()
mock_result2.fetchone.return_value = ('max_connections', '300')
mock_result2.fetchone.return_value = ("max_connections", "300")
def sql_side_effect(query):
if "SHOW PROCESSLIST" in query:
return mock_result1
elif "SHOW VARIABLES LIKE 'max_connections'" in query:
return mock_result2
return Mock()
mock_client.perform_raw_text_sql.side_effect = sql_side_effect
result = conn._get_connection_pool_stats()
assert "active_connections" in result
assert "max_connections" in result
assert result["active_connections"] >= 0
def test_get_slow_query_count(self):
"""Test retrieval of slow query count."""
# Create mock connection with actual methods
@@ -379,16 +356,16 @@ class TestOBConnectionPerformanceMetrics:
mock_client = Mock()
conn.client = mock_client
conn.logger = Mock()
mock_result = Mock()
mock_result.fetchone.return_value = (5,)
mock_client.perform_raw_text_sql.return_value = mock_result
result = conn._get_slow_query_count(threshold_seconds=1)
assert isinstance(result, int)
assert result >= 0
def test_estimate_qps(self):
"""Test QPS estimation."""
# Create mock connection with actual methods
@@ -396,17 +373,16 @@ class TestOBConnectionPerformanceMetrics:
mock_client = Mock()
conn.client = mock_client
conn.logger = Mock()
mock_result = Mock()
mock_result.fetchone.return_value = (10,)
mock_client.perform_raw_text_sql.return_value = mock_result
result = conn._estimate_qps()
assert isinstance(result, int)
assert result >= 0
if __name__ == "__main__":
pytest.main([__file__, "-v"])