gmakstutis/LightRAG
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Enhance workspace isolation test with distinct mock data and persistence

Browse source 
• Use different mock LLM per workspace
• Add persistent test directory
• Create workspace-specific responses
• Skip cleanup for inspection
This commit is contained in:
yangdx 2025-11-18 00:38:31 +08:00
parent b7b8d15632
commit 99262adaaa
1 changed files with 33 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
tests/test_workspace_isolation.py
View file

@ -813,21 +813,34 @@ async def test_lightrag_end_to_end_workspace_isolation():
print("=" * 60) print("=" * 60)
# Create temporary test directory # 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}") print(f"\n Using test directory: {test_dir}")
try: try:
# Mock LLM function # Factory function to create different mock LLM functions for each workspace
async def mock_llm_func( def create_mock_llm_func(workspace_name):
prompt, system_prompt=None, history_messages=[], **kwargs """Create a mock LLM function that returns different content based on workspace"""
) -> str: async def mock_llm_func(
# Return a mock response that simulates entity extraction in the correct format prompt, system_prompt=None, history_messages=[], **kwargs
# Format: entity<|#|>entity_name<|#|>entity_type<|#|>entity_description ) -> str:
# Format: relation<|#|>source_entity<|#|>target_entity<|#|>keywords<|#|>description # Return different responses based on workspace
return """entity<|#|>Artificial Intelligence<|#|>concept<|#|>AI is a field of computer science focused on creating intelligent machines. # 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. 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. relation<|#|>Machine Learning<|#|>Artificial Intelligence<|#|>subset, related field<|#|>Machine Learning is a key component and subset of Artificial Intelligence.
<|COMPLETE|>""" <|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 # Mock embedding function
async def mock_embedding_func(texts: list[str]) -> np.ndarray: 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 import LightRAG
from lightrag.utils import EmbeddingFunc 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( rag1 = LightRAG(
working_dir=test_dir, working_dir=test_dir,
workspace="project_a", workspace="project_a",
llm_model_func=mock_llm_func, llm_model_func=mock_llm_func_a,
embedding_func=EmbeddingFunc( embedding_func=EmbeddingFunc(
embedding_dim=384, embedding_dim=384,
max_token_size=8192, max_token_size=8192,
@ -853,7 +870,7 @@ relation<|#|>Machine Learning<|#|>Artificial Intelligence<|#|>subset, related fi
rag2 = LightRAG( rag2 = LightRAG(
working_dir=test_dir, working_dir=test_dir,
workspace="project_b", workspace="project_b",
llm_model_func=mock_llm_func, llm_model_func=mock_llm_func_b,
embedding_func=EmbeddingFunc( embedding_func=EmbeddingFunc(
embedding_dim=384, embedding_dim=384,
max_token_size=8192, max_token_size=8192,
@ -982,6 +999,8 @@ relation<|#|>Machine Learning<|#|>Artificial Intelligence<|#|>subset, related fi
finally: finally:
# Cleanup test directory # Cleanup test directory
if os.path.exists(test_dir): # if os.path.exists(test_dir):
shutil.rmtree(test_dir) # shutil.rmtree(test_dir)
print(f"\n Cleaned up test directory: {test_dir}") # print(f"\n Cleaned up test directory: {test_dir}")
print("Keep test directory for manual inspection:")
print(f" {test_dir}")