test(kb): add kb manager resilience tests

cover initialization failure and recovery scenarios to guard
against regressions in kb error handling

include reference assets under refs for test validation

Author: RC-CHN

tests/unit/test_kb_manager_resilience.py

@@ -0,0 +1,305 @@
+"""
+Unit tests for knowledge base manager resilience behavior.
+
+Tests the following scenarios:
+1. update_kb preserves old instance when re-initialization fails
+2. update_kb switches instance only after new instance initializes successfully
+3. _ensure_vec_db clears stale init_error after successful initialization
+
+These tests use lazy imports and mocks to avoid circular import issues
+in the astrbot core module chain.
+"""
+
+import sys
+import types
+from pathlib import Path
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+
+@pytest.fixture
+def stub_provider_manager_module():
+    """Stub provider manager module to avoid circular imports in unit tests."""
+    original_module = sys.modules.get("astrbot.core.provider.manager")
+    stub_module = types.ModuleType("astrbot.core.provider.manager")
+
+    class ProviderManager: ...
+
+    setattr(stub_module, "ProviderManager", ProviderManager)
+    sys.modules["astrbot.core.provider.manager"] = stub_module
+
+    try:
+        yield
+    finally:
+        if original_module is not None:
+            sys.modules["astrbot.core.provider.manager"] = original_module
+        else:
+            sys.modules.pop("astrbot.core.provider.manager", None)
+
+
+@pytest.fixture
+def mock_provider_manager():
+    """Create a mock ProviderManager."""
+    manager = MagicMock()
+    manager.get_provider_by_id = AsyncMock()
+    manager.acm = MagicMock()
+    manager.acm.default_conf = {}
+    return manager
+
+
+@pytest.fixture
+def mock_kb_db():
+    """Create a mock KBSQLiteDatabase."""
+    db = MagicMock()
+    db.get_db = MagicMock()
+    db.list_kbs = AsyncMock(return_value=[])
+    db.get_kb_by_id = AsyncMock()
+    return db
+
+
+@pytest.fixture
+def mock_knowledge_base():
+    """Create a mock KnowledgeBase instance."""
+    # Use lazy import to avoid circular import
+    from astrbot.core.knowledge_base.models import KnowledgeBase
+
+    kb = KnowledgeBase(
+        kb_name="test_kb",
+        description="Test knowledge base",
+        emoji="📚",
+        embedding_provider_id="test-embedding-provider",
+        rerank_provider_id=None,
+        chunk_size=512,
+        chunk_overlap=50,
+        top_k_dense=50,
+        top_k_sparse=50,
+        top_m_final=5,
+    )
+    return kb
+
+
+@pytest.fixture
+def mock_embedding_provider():
+    """Create a mock EmbeddingProvider."""
+    provider = MagicMock()
+    provider.get_embeddings_batch = AsyncMock(return_value=[[0.1, 0.2, 0.3]])
+    return provider
+
+
+@pytest.mark.asyncio
+async def test_update_kb_preserves_old_instance_when_reinit_fails(
+    stub_provider_manager_module,
+    mock_provider_manager,
+    mock_kb_db,
+    mock_knowledge_base,
+    mock_embedding_provider,
+):
+    """
+    Test that update_kb preserves the old KBHelper instance when
+    re-initialization fails, ensuring the knowledge base remains available.
+    """
+    # Lazy import to avoid circular import
+    from astrbot.core.knowledge_base.kb_helper import KBHelper
+    from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager
+
+    # Setup: create an existing KBHelper with working vec_db
+    mock_provider_manager.get_provider_by_id.return_value = mock_embedding_provider
+
+    # Create KBHelper using __new__ to avoid __init__ side effects
+    old_helper = KBHelper.__new__(KBHelper)
+    old_helper.kb = mock_knowledge_base
+    old_helper.prov_mgr = mock_provider_manager
+    old_helper.kb_db = mock_kb_db
+    old_helper.kb_root_dir = "/tmp/test_kb"
+    old_helper.chunker = MagicMock()
+    old_helper.init_error = None
+    old_helper.vec_db = MagicMock()  # Simulate existing working vec_db
+    old_helper.terminate = AsyncMock()
+
+    # Create KBManager and inject the existing helper
+    kb_mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager)
+    kb_mgr.provider_manager = mock_provider_manager
+    kb_mgr.kb_db = mock_kb_db
+    kb_mgr.kb_insts = {mock_knowledge_base.kb_id: old_helper}
+    kb_mgr.retrieval_manager = MagicMock()
+
+    # Mock KBHelper creation to simulate initialization failure
+    with patch.object(KBHelper, "initialize", new_callable=AsyncMock) as mock_init:
+        # First call (for new_helper) should fail
+        mock_init.side_effect = Exception("Embedding provider unavailable")
+
+        # Execute update_kb with a different embedding provider
+        result = await kb_mgr.update_kb(
+            kb_id=mock_knowledge_base.kb_id,
+            kb_name="updated_kb",
+            embedding_provider_id="new-embedding-provider",
+        )
+
+        # Verify: the old helper should be returned, not a new one
+        assert result is not None
+        assert result is old_helper
+        assert kb_mgr.kb_insts[mock_knowledge_base.kb_id] is old_helper
+
+        # Verify: old helper's vec_db should still be available
+        assert hasattr(result, "vec_db")
+        assert result.vec_db is not None
+
+        # Verify: failure does not replace the existing helper state
+        assert result.init_error is None
+        assert result.kb.kb_name == "test_kb"
+        assert result.kb.embedding_provider_id == "test-embedding-provider"
+
+
+@pytest.mark.asyncio
+async def test_update_kb_switches_instance_only_after_new_reinit_success(
+    stub_provider_manager_module,
+    mock_provider_manager,
+    mock_kb_db,
+    mock_knowledge_base,
+    mock_embedding_provider,
+):
+    """
+    Test that update_kb only switches to the new KBHelper instance
+    after the new instance successfully initializes.
+    """
+    # Lazy import to avoid circular import
+    from astrbot.core.knowledge_base.kb_helper import KBHelper
+    from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager
+
+    # Setup: create an existing KBHelper
+    mock_provider_manager.get_provider_by_id.return_value = mock_embedding_provider
+
+    old_helper = KBHelper.__new__(KBHelper)
+    old_helper.kb = mock_knowledge_base
+    old_helper.prov_mgr = mock_provider_manager
+    old_helper.kb_db = mock_kb_db
+    old_helper.kb_root_dir = "/tmp/test_kb"
+    old_helper.chunker = MagicMock()
+    old_helper.init_error = None
+    old_helper.vec_db = MagicMock()
+    old_helper.terminate = AsyncMock()
+
+    kb_mgr = KnowledgeBaseManager.__new__(KnowledgeBaseManager)
+    kb_mgr.provider_manager = mock_provider_manager
+    kb_mgr.kb_db = mock_kb_db
+    kb_mgr.kb_insts = {mock_knowledge_base.kb_id: old_helper}
+    kb_mgr.retrieval_manager = MagicMock()
+
+    # Mock session context for database operations
+    mock_session = MagicMock()
+    mock_session.add = MagicMock()
+    mock_session.commit = AsyncMock()
+    mock_session.refresh = AsyncMock()
+
+    mock_db_context = MagicMock()
+    mock_db_context.__aenter__ = AsyncMock(return_value=mock_session)
+    mock_db_context.__aexit__ = AsyncMock()
+    mock_kb_db.get_db.return_value = mock_db_context
+
+    # Mock KBHelper.initialize to succeed
+    with patch.object(KBHelper, "initialize", new_callable=AsyncMock) as mock_init:
+        mock_init.return_value = None
+
+        # Execute update_kb
+        result = await kb_mgr.update_kb(
+            kb_id=mock_knowledge_base.kb_id,
+            kb_name="updated_kb",
+            embedding_provider_id="new-embedding-provider",
+        )
+
+        # Verify: a new helper should be returned
+        assert result is not None
+        assert result is not old_helper
+        assert result.init_error is None
+        assert kb_mgr.kb_insts[mock_knowledge_base.kb_id] is result
+
+        # Verify: old helper should be terminated
+        old_helper.terminate.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_ensure_vec_db_clears_stale_init_error(
+    stub_provider_manager_module,
+    mock_provider_manager,
+    mock_kb_db,
+    mock_knowledge_base,
+    mock_embedding_provider,
+):
+    """
+    Test that _ensure_vec_db clears the init_error attribute
+    after successful initialization, removing stale error state.
+    """
+    # Lazy import to avoid circular import
+    from astrbot.core.knowledge_base.kb_helper import KBHelper
+
+    # Setup: create KBHelper with stale init_error
+    mock_provider_manager.get_provider_by_id.return_value = mock_embedding_provider
+
+    helper = KBHelper.__new__(KBHelper)
+    helper.kb = mock_knowledge_base
+    helper.prov_mgr = mock_provider_manager
+    helper.kb_db = mock_kb_db
+    helper.kb_root_dir = "/tmp/test_kb"
+    helper.chunker = MagicMock()
+    helper.init_error = "Previous initialization failed"
+    helper.kb_dir = Path("/tmp/test_kb") / mock_knowledge_base.kb_id
+    helper.kb_medias_dir = helper.kb_dir / "medias" / mock_knowledge_base.kb_id
+    helper.kb_files_dir = helper.kb_dir / "files" / mock_knowledge_base.kb_id
+
+    # Mock FaissVecDB initialization
+    mock_vec_db = MagicMock()
+    mock_vec_db.initialize = AsyncMock()
+    mock_vec_db.close = AsyncMock()
+
+    with patch(
+        "astrbot.core.knowledge_base.kb_helper.FaissVecDB",
+        return_value=mock_vec_db,
+    ):
+        # Execute _ensure_vec_db
+        await helper._ensure_vec_db()
+
+        # Verify: init_error should be cleared
+        assert helper.init_error is None
+        assert helper.vec_db is mock_vec_db
+
+
+@pytest.mark.asyncio
+async def test_ensure_vec_db_sets_init_error_on_failure(
+    stub_provider_manager_module,
+    mock_provider_manager,
+    mock_kb_db,
+    mock_knowledge_base,
+):
+    """
+    Test that _ensure_vec_db does NOT clear init_error when
+    initialization fails, preserving the error state.
+    """
+    # Lazy import to avoid circular import
+    from astrbot.core.knowledge_base.kb_helper import KBHelper
+
+    # Setup: provider unavailable
+    mock_provider_manager.get_provider_by_id.return_value = None
+
+    helper = KBHelper.__new__(KBHelper)
+    helper.kb = mock_knowledge_base
+    helper.prov_mgr = mock_provider_manager
+    helper.kb_db = mock_kb_db
+    helper.kb_root_dir = "/tmp/test_kb"
+    helper.chunker = MagicMock()
+    helper.init_error = "Previous initialization failed"
+    helper.kb_dir = Path("/tmp/test_kb") / mock_knowledge_base.kb_id
+    helper.kb_medias_dir = helper.kb_dir / "medias" / mock_knowledge_base.kb_id
+    helper.kb_files_dir = helper.kb_dir / "files" / mock_knowledge_base.kb_id
+
+    # Execute _ensure_vec_db - should raise exception
+    try:
+        await helper._ensure_vec_db()
+        pytest.fail("Expected exception but none was raised")
+    except ValueError as e:
+        # Verify: exception should be raised
+        assert "无法找到" in str(e) or "未配置" in str(e)
+
+        # Verify: init_error should NOT be cleared (still has previous error)
+        # Note: _ensure_vec_db doesn't set init_error; that's done by the caller
+        assert helper.init_error is not None