From 99262adaaa54e9d3c29cc14718f02c5fc3a8a91c Mon Sep 17 00:00:00 2001 From: yangdx Date: Tue, 18 Nov 2025 00:38:31 +0800 Subject: [PATCH] Enhance workspace isolation test with distinct mock data and persistence MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit • Use different mock LLM per workspace • Add persistent test directory • Create workspace-specific responses • Skip cleanup for inspection --- tests/test_workspace_isolation.py | 47 ++++++++++++++++++++++--------- 1 file changed, 33 insertions(+), 14 deletions(-) diff --git a/tests/test_workspace_isolation.py b/tests/test_workspace_isolation.py index 420e4d9b..4a6f284b 100644 --- a/tests/test_workspace_isolation.py +++ b/tests/test_workspace_isolation.py @@ -813,21 +813,34 @@ async def test_lightrag_end_to_end_workspace_isolation(): print("=" * 60) # Create temporary test directory - test_dir = tempfile.mkdtemp(prefix="lightrag_test_e2e_") + # test_dir = tempfile.mkdtemp(prefix="lightrag_test_e2e_") + test_dir = str(Path(__file__).parent.parent / "temp/e2e_workspace_isolation") + if os.path.exists(test_dir): + shutil.rmtree(test_dir) + os.makedirs(test_dir, exist_ok=True) print(f"\n Using test directory: {test_dir}") try: - # Mock LLM function - async def mock_llm_func( - prompt, system_prompt=None, history_messages=[], **kwargs - ) -> str: - # Return a mock response that simulates entity extraction in the correct format - # Format: entity<|#|>entity_name<|#|>entity_type<|#|>entity_description - # Format: relation<|#|>source_entity<|#|>target_entity<|#|>keywords<|#|>description - return """entity<|#|>Artificial Intelligence<|#|>concept<|#|>AI is a field of computer science focused on creating intelligent machines. + # Factory function to create different mock LLM functions for each workspace + def create_mock_llm_func(workspace_name): + """Create a mock LLM function that returns different content based on workspace""" + async def mock_llm_func( + prompt, system_prompt=None, history_messages=[], **kwargs + ) -> str: + # Return different responses based on workspace + # Format: entity<|#|>entity_name<|#|>entity_type<|#|>entity_description + # Format: relation<|#|>source_entity<|#|>target_entity<|#|>keywords<|#|>description + if workspace_name == "project_a": + return """entity<|#|>Artificial Intelligence<|#|>concept<|#|>AI is a field of computer science focused on creating intelligent machines. entity<|#|>Machine Learning<|#|>concept<|#|>Machine Learning is a subset of AI that enables systems to learn from data. relation<|#|>Machine Learning<|#|>Artificial Intelligence<|#|>subset, related field<|#|>Machine Learning is a key component and subset of Artificial Intelligence. <|COMPLETE|>""" + else: # project_b + return """entity<|#|>Deep Learning<|#|>concept<|#|>Deep Learning is a subset of machine learning using neural networks with multiple layers. +entity<|#|>Neural Networks<|#|>concept<|#|>Neural Networks are computing systems inspired by biological neural networks. +relation<|#|>Deep Learning<|#|>Neural Networks<|#|>uses, composed of<|#|>Deep Learning uses multiple layers of Neural Networks to learn representations. +<|COMPLETE|>""" + return mock_llm_func # Mock embedding function async def mock_embedding_func(texts: list[str]) -> np.ndarray: @@ -839,10 +852,14 @@ relation<|#|>Machine Learning<|#|>Artificial Intelligence<|#|>subset, related fi from lightrag import LightRAG from lightrag.utils import EmbeddingFunc + # Create different mock LLM functions for each workspace + mock_llm_func_a = create_mock_llm_func("project_a") + mock_llm_func_b = create_mock_llm_func("project_b") + rag1 = LightRAG( working_dir=test_dir, workspace="project_a", - llm_model_func=mock_llm_func, + llm_model_func=mock_llm_func_a, embedding_func=EmbeddingFunc( embedding_dim=384, max_token_size=8192, @@ -853,7 +870,7 @@ relation<|#|>Machine Learning<|#|>Artificial Intelligence<|#|>subset, related fi rag2 = LightRAG( working_dir=test_dir, workspace="project_b", - llm_model_func=mock_llm_func, + llm_model_func=mock_llm_func_b, embedding_func=EmbeddingFunc( embedding_dim=384, max_token_size=8192, @@ -982,6 +999,8 @@ relation<|#|>Machine Learning<|#|>Artificial Intelligence<|#|>subset, related fi finally: # Cleanup test directory - if os.path.exists(test_dir): - shutil.rmtree(test_dir) - print(f"\n Cleaned up test directory: {test_dir}") + # if os.path.exists(test_dir): + # shutil.rmtree(test_dir) + # print(f"\n Cleaned up test directory: {test_dir}") + print("Keep test directory for manual inspection:") + print(f" {test_dir}")