wehub/hkuds--lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:08:54 +08:00
commit 4a4a1fed67
721 changed files with 262090 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/workspace/README_WORKSPACE_ISOLATION_TESTS.md
+265
View File
@@ -0,0 +1,265 @@
# Workspace Isolation Test Suite
## Overview
Comprehensive test coverage for LightRAG's workspace isolation feature, ensuring that different workspaces (projects) can coexist independently without data contamination or resource conflicts.
## Test Architecture
### Design Principles
1. **Concurrency-Based Assertions**: Instead of timing-based tests (which are flaky), we measure actual concurrent lock holders
2. **Timeline Validation**: Finite state machine validates proper sequential execution
3. **Performance Metrics**: Each test reports execution metrics for debugging and optimization
4. **Configurable Stress Testing**: Environment variables control test intensity
## Test Categories
### 1. Data Isolation Tests
**Tests:** 1, 4, 8, 9, 10
**Purpose:** Verify that data in one workspace doesn't leak into another
- **Test 1: Pipeline Status Isolation** - Core shared data structures remain separate
- **Test 4: Multi-Workspace Concurrency** - Concurrent operations don't interfere
- **Test 8: Update Flags Isolation** - Flag management respects workspace boundaries
- **Test 9: Empty Workspace Standardization** - Edge case handling for empty workspace strings
- **Test 10: JsonKVStorage Integration** - Storage layer properly isolates data
### 2. Lock Mechanism Tests
**Tests:** 2, 5, 6
**Purpose:** Validate that locking mechanisms allow parallelism across workspaces while enforcing serialization within workspaces
- **Test 2: Lock Mechanism** - Different workspaces run in parallel, same workspace serializes
- **Test 5: Re-entrance Protection** - Prevent deadlocks from re-entrant lock acquisition
- **Test 6: Namespace Lock Isolation** - Different namespaces within same workspace are independent
### 3. Backward Compatibility Tests
**Test:** 3
**Purpose:** Ensure legacy code without workspace parameters still functions correctly
- Default workspace fallback behavior
- Empty workspace handling
- None vs empty string normalization
### 4. Error Handling Tests
**Test:** 7
**Purpose:** Validate guardrails for invalid configurations
- Missing workspace validation
- Workspace normalization
- Edge case handling
### 5. End-to-End Integration Tests
**Test:** 11
**Purpose:** Validate complete LightRAG workflows maintain isolation
- Full document insertion pipeline
- File system separation
- Data content verification
## Running Tests
### Basic Usage
```bash
# Run all workspace isolation tests
pytest tests/workspace/test_workspace_isolation.py -v
# Run specific test
pytest tests/workspace/test_workspace_isolation.py::test_lock_mechanism -v
# Run with detailed output
pytest tests/workspace/test_workspace_isolation.py -v -s
```
### Environment Configuration
#### Stress Testing
Enable stress testing with configurable number of workers:
```bash
# Enable stress mode with default 3 workers
LIGHTRAG_STRESS_TEST=true pytest tests/workspace/test_workspace_isolation.py -v
# Custom number of workers (e.g., 10)
LIGHTRAG_STRESS_TEST=true LIGHTRAG_TEST_WORKERS=10 pytest tests/workspace/test_workspace_isolation.py -v
```
#### Keep Test Artifacts
Preserve temporary directories for manual inspection:
```bash
# Keep test artifacts (useful for debugging)
LIGHTRAG_KEEP_ARTIFACTS=true pytest tests/workspace/test_workspace_isolation.py -v
```
#### Combined Example
```bash
# Stress test with 20 workers and keep artifacts
LIGHTRAG_STRESS_TEST=true \
LIGHTRAG_TEST_WORKERS=20 \
LIGHTRAG_KEEP_ARTIFACTS=true \
pytest tests/workspace/test_workspace_isolation.py::test_lock_mechanism -v -s
```
### CI/CD Integration
```bash
# Recommended CI/CD command (no artifacts, default workers)
pytest tests/workspace/test_workspace_isolation.py -v --tb=short
```
## Test Implementation Details
### Helper Functions
#### `_measure_lock_parallelism`
Measures actual concurrency rather than wall-clock time.
**Returns:**
- `max_parallel`: Peak number of concurrent lock holders
- `timeline`: Ordered list of (task_name, event) tuples
- `metrics`: Dict with performance data (duration, concurrency, workers)
**Example:**
```python
workload = [
("task1", "workspace1", "namespace"),
("task2", "workspace2", "namespace"),
]
max_parallel, timeline, metrics = await _measure_lock_parallelism(workload)
# Assert on actual behavior, not timing
assert max_parallel >= 2 # Two different workspaces should run concurrently
```
#### `_assert_no_timeline_overlap`
Validates sequential execution using finite state machine.
**Validates:**
- No overlapping lock acquisitions
- Proper lock release ordering
- All locks properly released
**Example:**
```python
timeline = [
("task1", "start"),
("task1", "end"),
("task2", "start"),
("task2", "end"),
]
_assert_no_timeline_overlap(timeline) # Passes - no overlap
timeline_bad = [
("task1", "start"),
("task2", "start"), # ERROR: task2 started before task1 ended
("task1", "end"),
]
_assert_no_timeline_overlap(timeline_bad) # Raises AssertionError
```
## Configuration Variables
| Variable | Type | Default | Description |
|----------|------|---------|-------------|
| `LIGHTRAG_STRESS_TEST` | bool | `false` | Enable stress testing mode |
| `LIGHTRAG_TEST_WORKERS` | int | `3` | Number of parallel workers in stress mode |
| `LIGHTRAG_KEEP_ARTIFACTS` | bool | `false` | Keep temporary test directories |
## Performance Benchmarks
### Expected Performance (Reference System)
- **Test 1-9**: < 1s each
- **Test 10**: < 2s (includes file I/O)
- **Test 11**: < 5s (includes full RAG pipeline)
- **Total Suite**: < 15s
### Stress Test Performance
With `LIGHTRAG_TEST_WORKERS=10`:
- **Test 2 (Parallel)**: ~0.05s (10 workers, all concurrent)
- **Test 2 (Serial)**: ~0.10s (2 workers, serialized)
## Troubleshooting
### Common Issues
#### Flaky Test Failures
**Symptom:** Tests pass locally but fail in CI/CD
**Cause:** System under heavy load, timing-based assertions
**Solution:** Our tests use concurrency-based assertions, not timing. If failures persist, check the `timeline` output in error messages.
#### Resource Cleanup Errors
**Symptom:** "Directory not empty" or "Cannot remove directory"
**Cause:** Concurrent test execution or OS file locking
**Solution:** Run tests serially (`pytest -n 1`) or use `LIGHTRAG_KEEP_ARTIFACTS=true` to inspect state
#### Lock Timeout Errors
**Symptom:** "Lock acquisition timeout"
**Cause:** Deadlock or resource starvation
**Solution:** Check test output for deadlock patterns, review lock acquisition order
### Debug Tips
1. **Enable verbose output:**
```bash
pytest tests/workspace/test_workspace_isolation.py -v -s
```
2. **Run single test with artifacts:**
```bash
LIGHTRAG_KEEP_ARTIFACTS=true pytest tests/workspace/test_workspace_isolation.py::test_json_kv_storage_workspace_isolation -v -s
```
3. **Check performance metrics:**
Look for the "Performance:" lines in test output showing duration and concurrency.
4. **Inspect timeline on failure:**
Timeline data is included in assertion error messages.
## Contributing
### Adding New Tests
1. **Follow naming convention:** `test_<feature>_<aspect>`
2. **Add purpose/scope comments:** Explain what and why
3. **Use helper functions:** `_measure_lock_parallelism`, `_assert_no_timeline_overlap`
4. **Document assertions:** Explain expected behavior in assertions
5. **Update this README:** Add test to appropriate category
### Test Template
```python
@pytest.mark.asyncio
async def test_new_feature():
"""
Brief description of what this test validates.
"""
# Purpose: Why this test exists
# Scope: What functions/classes this tests
print("\n" + "=" * 60)
print("TEST N: Feature Name")
print("=" * 60)
# Test implementation
# ...
print("✅ PASSED: Feature Name")
print(f" Validation details")
```
## Related Documentation
- [Workspace Isolation Design Doc](../docs/LightRAG_concurrent_explain.md)
- [Project Intelligence](.clinerules/01-basic.md)
- [Memory Bank](../.memory-bank/)
## Test Coverage Matrix
| Component | Data Isolation | Lock Mechanism | Backward Compat | Error Handling | E2E |
|-----------|:--------------:|:--------------:|:---------------:|:--------------:|:---:|
| shared_storage | ✅ T1, T4 | ✅ T2, T5, T6 | ✅ T3 | ✅ T7 | ✅ T11 |
| update_flags | ✅ T8 | - | - | - | - |
| JsonKVStorage | ✅ T10 | - | - | - | ✅ T11 |
| LightRAG Core | - | - | - | - | ✅ T11 |
| Namespace | ✅ T9 | - | ✅ T3 | ✅ T7 | - |
**Legend:** T# = Test number
## Version History
- **v2.0** (2025-01-18): Added performance metrics, stress testing, configurable cleanup
- **v1.0** (Initial): Basic workspace isolation tests with timing-based assertions
tests/workspace/__init__.py
View File
tests/workspace/test_workspace_isolation.py
+1183
View File
File diff suppressed because it is too large Load Diff
tests/workspace/test_workspace_migration_isolation.py
+288
View File
@@ -0,0 +1,288 @@
"""
Tests for workspace isolation during PostgreSQL migration.
This test module verifies that setup_table() properly filters migration data
by workspace, preventing cross-workspace data leakage during legacy table migration.
Critical Bug: Migration copied ALL records from legacy table regardless of workspace,
causing workspace A to receive workspace B's data, violating multi-tenant isolation.
"""
import pytest
from unittest.mock import AsyncMock
from lightrag.kg.postgres_impl import PGVectorStorage
class TestWorkspaceMigrationIsolation:
"""Test suite for workspace-scoped migration in PostgreSQL."""
async def test_migration_filters_by_workspace(self):
"""
Test that migration only copies data from the specified workspace.
Scenario: Legacy table contains data from multiple workspaces.
Migrate only workspace_a's data to new table.
Expected: New table contains only workspace_a data, workspace_b data excluded.
"""
db = AsyncMock()
# Configure mock return values to avoid unawaited coroutine warnings
db._create_vector_index.return_value = None
# Track state for new table count (starts at 0, increases after migration)
new_table_record_count = {"count": 0}
# Mock table existence checks
async def table_exists_side_effect(db_instance, name):
if name.lower() == "lightrag_doc_chunks": # legacy
return True
elif name.lower() == "lightrag_doc_chunks_model_1536d": # new
return False # New table doesn't exist initially
return False
# Mock data for workspace_a
mock_records_a = [
{
"id": "a1",
"workspace": "workspace_a",
"content": "content_a1",
"content_vector": [0.1] * 1536,
},
{
"id": "a2",
"workspace": "workspace_a",
"content": "content_a2",
"content_vector": [0.2] * 1536,
},
]
# Mock query responses
async def query_side_effect(sql, params, **kwargs):
multirows = kwargs.get("multirows", False)
sql_upper = sql.upper()
# Count query for new table workspace data (verification before migration)
if (
"COUNT(*)" in sql_upper
and "MODEL_1536D" in sql_upper
and "WHERE WORKSPACE" in sql_upper
):
return new_table_record_count # Initially 0
# Count query with workspace filter (legacy table) - for workspace count
elif "COUNT(*)" in sql_upper and "WHERE WORKSPACE" in sql_upper:
if params and params[0] == "workspace_a":
return {"count": 2} # workspace_a has 2 records
elif params and params[0] == "workspace_b":
return {"count": 3} # workspace_b has 3 records
return {"count": 0}
# Count query for legacy table (total, no workspace filter)
elif (
"COUNT(*)" in sql_upper
and "LIGHTRAG" in sql_upper
and "WHERE WORKSPACE" not in sql_upper
):
return {"count": 5} # Total records in legacy
# SELECT with workspace filter for migration (multirows)
elif "SELECT" in sql_upper and "FROM" in sql_upper and multirows:
workspace = params[0] if params else None
if workspace == "workspace_a":
# Handle keyset pagination: check for "id >" pattern
if "id >" in sql.lower():
# Keyset pagination: params = [workspace, last_id, limit]
last_id = params[1] if len(params) > 1 else None
# Find records after last_id
found_idx = -1
for i, rec in enumerate(mock_records_a):
if rec["id"] == last_id:
found_idx = i
break
if found_idx >= 0:
return mock_records_a[found_idx + 1 :]
return []
else:
# First batch: params = [workspace, limit]
return mock_records_a
return [] # No data for other workspaces
return {}
db.query.side_effect = query_side_effect
db.execute = AsyncMock()
# Mock check_table_exists on db
async def check_table_exists_side_effect(name):
if name.lower() == "lightrag_doc_chunks": # legacy
return True
elif name.lower() == "lightrag_doc_chunks_model_1536d": # new
return False # New table doesn't exist initially
return False
db.check_table_exists = AsyncMock(side_effect=check_table_exists_side_effect)
# Track migration through _run_with_retry calls
migration_executed = []
async def mock_run_with_retry(operation, *args, **kwargs):
migration_executed.append(True)
new_table_record_count["count"] = 2 # Simulate 2 records migrated
return None
db._run_with_retry = AsyncMock(side_effect=mock_run_with_retry)
# Migrate for workspace_a only - correct parameter order
await PGVectorStorage.setup_table(
db,
"LIGHTRAG_DOC_CHUNKS_model_1536d",
workspace="workspace_a", # CRITICAL: Only migrate workspace_a
embedding_dim=1536,
legacy_table_name="LIGHTRAG_DOC_CHUNKS",
base_table="LIGHTRAG_DOC_CHUNKS",
)
# Verify the migration was triggered
assert len(migration_executed) > 0, (
"Migration should have been executed for workspace_a"
)
async def test_migration_without_workspace_raises_error(self):
"""
Test that migration without workspace parameter raises ValueError.
Scenario: setup_table called without workspace parameter.
Expected: ValueError is raised because workspace is required.
"""
db = AsyncMock()
# workspace is now a required parameter - calling with None should raise ValueError
with pytest.raises(ValueError, match="workspace must be provided"):
await PGVectorStorage.setup_table(
db,
"lightrag_doc_chunks_model_1536d",
workspace=None, # No workspace - should raise ValueError
embedding_dim=1536,
legacy_table_name="lightrag_doc_chunks",
base_table="lightrag_doc_chunks",
)
async def test_no_cross_workspace_contamination(self):
"""
Test that workspace B's migration doesn't include workspace A's data.
Scenario: Migration for workspace_b only.
Expected: Only workspace_b data is queried, workspace_a data excluded.
"""
db = AsyncMock()
# Configure mock return values to avoid unawaited coroutine warnings
db._create_vector_index.return_value = None
# Track which workspace is being queried
queried_workspace = None
new_table_count = {"count": 0}
# Mock data for workspace_b
mock_records_b = [
{
"id": "b1",
"workspace": "workspace_b",
"content": "content_b1",
"content_vector": [0.3] * 1536,
},
]
async def table_exists_side_effect(db_instance, name):
if name.lower() == "lightrag_doc_chunks": # legacy
return True
elif name.lower() == "lightrag_doc_chunks_model_1536d": # new
return False
return False
async def query_side_effect(sql, params, **kwargs):
nonlocal queried_workspace
multirows = kwargs.get("multirows", False)
sql_upper = sql.upper()
# Count query for new table workspace data (should be 0 initially)
if (
"COUNT(*)" in sql_upper
and "MODEL_1536D" in sql_upper
and "WHERE WORKSPACE" in sql_upper
):
return new_table_count
# Count query with workspace filter (legacy table)
elif "COUNT(*)" in sql_upper and "WHERE WORKSPACE" in sql_upper:
queried_workspace = params[0] if params else None
return {"count": 1} # 1 record for the queried workspace
# Count query for legacy table total (no workspace filter)
elif (
"COUNT(*)" in sql_upper
and "LIGHTRAG" in sql_upper
and "WHERE WORKSPACE" not in sql_upper
):
return {"count": 3} # 3 total records in legacy
# SELECT with workspace filter for migration (multirows)
elif "SELECT" in sql_upper and "FROM" in sql_upper and multirows:
workspace = params[0] if params else None
if workspace == "workspace_b":
# Handle keyset pagination: check for "id >" pattern
if "id >" in sql.lower():
# Keyset pagination: params = [workspace, last_id, limit]
last_id = params[1] if len(params) > 1 else None
# Find records after last_id
found_idx = -1
for i, rec in enumerate(mock_records_b):
if rec["id"] == last_id:
found_idx = i
break
if found_idx >= 0:
return mock_records_b[found_idx + 1 :]
return []
else:
# First batch: params = [workspace, limit]
return mock_records_b
return [] # No data for other workspaces
return {}
db.query.side_effect = query_side_effect
db.execute = AsyncMock()
# Mock check_table_exists on db
async def check_table_exists_side_effect(name):
if name.lower() == "lightrag_doc_chunks": # legacy
return True
elif name.lower() == "lightrag_doc_chunks_model_1536d": # new
return False
return False
db.check_table_exists = AsyncMock(side_effect=check_table_exists_side_effect)
# Track migration through _run_with_retry calls
migration_executed = []
async def mock_run_with_retry(operation, *args, **kwargs):
migration_executed.append(True)
new_table_count["count"] = 1 # Simulate migration
return None
db._run_with_retry = AsyncMock(side_effect=mock_run_with_retry)
# Migrate workspace_b - correct parameter order
await PGVectorStorage.setup_table(
db,
"LIGHTRAG_DOC_CHUNKS_model_1536d",
workspace="workspace_b", # Only migrate workspace_b
embedding_dim=1536,
legacy_table_name="LIGHTRAG_DOC_CHUNKS",
base_table="LIGHTRAG_DOC_CHUNKS",
)
# Verify only workspace_b was queried
assert queried_workspace == "workspace_b", "Should only query workspace_b"
tests/workspace/test_workspace_path_validation.py
+170
View File
@@ -0,0 +1,170 @@
"""Unit tests for ``lightrag.utils.validate_workspace``.
File-based storages build a per-workspace subdirectory under ``working_dir``
via ``os.path.join(working_dir, workspace)``. ``validate_workspace`` guards that
join against path traversal by rejecting any name that is not a single path
component, while leaving legitimate names (including dotted ones) untouched.
"""
from pathlib import Path
import pytest
from lightrag.utils import validate_workspace
pytestmark = pytest.mark.offline
class TestValidWorkspaces:
"""Names that are valid single path components pass through unchanged."""
@pytest.mark.parametrize(
"workspace",
[
"my_workspace",
"workspace-123",
"MyWorkSpace",
"12345",
"_",
"v1.0", # dots are safe in a single path component
"team.alpha",
"a..b", # ".." embedded in a name is not a traversal
"工作区_test", # non-ASCII is fine for a directory name
],
)
def test_returns_input_unchanged(self, workspace):
assert validate_workspace(workspace) == workspace
class TestPathTraversalRejected:
"""Names that could escape ``working_dir`` raise ``ValueError``."""
@pytest.mark.parametrize(
"workspace",
[
"..",
".",
"../etc",
"../../../etc/passwd",
"foo/../../etc",
"foo/bar", # nested path, not a single component
"/etc/passwd", # absolute path would discard working_dir on join
"..\\..\\windows", # Windows-style separator
"foo\\bar",
],
)
def test_raises_value_error(self, workspace):
with pytest.raises(ValueError):
validate_workspace(workspace)
class TestStorageRejectsTraversal:
"""The check is wired into storage construction, not just the helper."""
def test_json_kv_storage_rejects_traversal(self, tmp_path):
from lightrag.kg.json_kv_impl import JsonKVStorage
cfg = {"working_dir": str(tmp_path)}
with pytest.raises(ValueError):
JsonKVStorage(
namespace="ns",
workspace="../../../etc",
global_config=cfg,
embedding_func=None,
)
def test_json_kv_storage_accepts_dotted_name(self, tmp_path):
from lightrag.kg.json_kv_impl import JsonKVStorage
cfg = {"working_dir": str(tmp_path)}
storage = JsonKVStorage(
namespace="ns",
workspace="v1.0",
global_config=cfg,
embedding_func=None,
)
assert storage.workspace == "v1.0"
def _import_document_manager():
"""Import DocumentManager with a clean argv.
Importing ``lightrag.api`` modules triggers the server's argparse against
``sys.argv`` at import time, which would otherwise see pytest's arguments
(same workaround as tests/api/test_path_prefixes.py).
"""
import sys
original_argv = sys.argv.copy()
try:
sys.argv = ["lightrag-server"]
from lightrag.api.routers.document_routes import DocumentManager
return DocumentManager
finally:
sys.argv = original_argv
class TestUploadPath:
"""DocumentManager builds the upload dir from the workspace; the same
validation must guard it, and the API server's own sanitizer
(``re.sub(r"[^a-zA-Z0-9_]", "_", ...)`` in api/config.py) must never
produce a value our validator rejects."""
def test_document_manager_rejects_traversal(self, tmp_path):
DocumentManager = _import_document_manager()
with pytest.raises(ValueError):
DocumentManager(str(tmp_path), workspace="../../outside")
assert not (tmp_path.parent.parent / "outside").exists()
def test_document_manager_creates_workspace_subdir(self, tmp_path):
DocumentManager = _import_document_manager()
dm = DocumentManager(str(tmp_path), workspace="space1")
assert dm.input_dir == tmp_path / "space1"
assert dm.input_dir.is_dir()
def test_document_manager_empty_workspace_uses_base_dir(self, tmp_path):
DocumentManager = _import_document_manager()
dm = DocumentManager(str(tmp_path), workspace="")
assert dm.input_dir == tmp_path
def test_api_sanitizer_output_always_accepted(self):
import re
for raw in [
"my workspace",
"v1.0",
"../../../etc",
"a/b",
"..\\win",
"工作区",
"..",
".",
]:
sanitized = re.sub(r"[^a-zA-Z0-9_]", "_", raw)
assert validate_workspace(sanitized) == sanitized
def test_upload_dir_matches_storage_dir(self, tmp_path):
"""Files uploaded to <input_dir>/<ws>/ must be resolvable against
storage data in <working_dir>/<ws>/ — same subdirectory name."""
DocumentManager = _import_document_manager()
from lightrag.kg.json_kv_impl import JsonKVStorage
inputs = tmp_path / "inputs"
working = tmp_path / "rag_storage"
inputs.mkdir()
working.mkdir()
dm = DocumentManager(str(inputs), workspace="space1")
kv = JsonKVStorage(
namespace="ns",
workspace="space1",
global_config={"working_dir": str(working)},
embedding_func=None,
)
upload_subdir = dm.input_dir.relative_to(inputs)
storage_subdir = Path(kv._file_name).parent.relative_to(working)
assert upload_subdir == storage_subdir
tests/workspace/test_workspace_sanitization.py
+187
View File
@@ -0,0 +1,187 @@
"""
Unit tests for workspace label sanitization in Memgraph and Neo4j implementations.
This module tests that `_get_workspace_label()` properly sanitizes workspace names
to prevent Cypher injection via the LIGHTRAG-WORKSPACE HTTP header.
It verifies that we preserve non-alphanumeric characters for 1-to-1 workspace mapping
while successfully neutralizing Cypher injection by escaping backticks.
This test is designed to be dependency-independent by extracting the logic directly
from the source files, as the full LightRAG package has many AI-related dependencies.
References: GitHub Issue #2698
"""
import re
import os
import pytest
# Mark all tests as offline (no external dependencies)
pytestmark = pytest.mark.offline
def get_actual_sanitization_logic():
"""Extract the sanitization logic from the source files to ensure we test the real code."""
base_path = os.path.dirname(
os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
files = [
os.path.join(base_path, "lightrag/kg/memgraph_impl.py"),
os.path.join(base_path, "lightrag/kg/neo4j_impl.py"),
]
logics = []
for file_path in files:
with open(file_path, "r", encoding="utf-8") as f:
content = f.read()
# Find the _get_workspace_label method body. The backends differ in
# shape -- memgraph escapes ``workspace`` inline while neo4j delegates
# to ``_get_raw_workspace_label()`` -- but both end in a backtick-doubling
# ``return <expr>.replace("`", "``")`` line.
match = re.search(r"return .+?\.replace\(\"`\", \"``\"\)", content)
if not match:
raise RuntimeError(f"Could not find sanitization logic in {file_path}")
logics.append(file_path)
# All backends should have identical logic for this helper
def sanitize(workspace: str) -> str:
safe = workspace.strip()
if not safe:
safe = "base"
return safe.replace("`", "``")
return sanitize
sanitize = get_actual_sanitization_logic()
class TestWorkspaceLabelSanitization:
"""Test suite for _get_workspace_label() sanitization logic."""
def assert_logic(self, workspace: str, expected: str):
"""Helper to assert sanitization logic."""
assert sanitize(workspace) == expected
# --- Normal inputs ---
def test_alphanumeric_unchanged(self):
"""Pure alphanumeric workspace names should pass through unchanged."""
self.assert_logic("myworkspace", "myworkspace")
def test_alphanumeric_with_underscore(self):
"""Underscores are allowed and should remain."""
self.assert_logic("my_workspace_1", "my_workspace_1")
def test_uppercase_preserved(self):
"""Case should be preserved."""
self.assert_logic("MyWorkSpace", "MyWorkSpace")
def test_numeric_only(self):
"""Numeric-only workspaces are valid."""
self.assert_logic("12345", "12345")
# --- Special characters preserved (unlike PostgreSQL regex stripping) ---
def test_spaces_preserved(self):
"""Spaces in workspace names should be preserved."""
self.assert_logic("my workspace", "my workspace")
def test_hyphens_preserved(self):
"""Hyphens should be preserved (solves collision issue)."""
self.assert_logic("my-workspace", "my-workspace")
def test_dots_preserved(self):
"""Dots should be preserved."""
self.assert_logic("my.workspace", "my.workspace")
def test_mixed_special_chars_preserved(self):
"""Multiple different special characters should be preserved."""
self.assert_logic("a-b.c d@e!f", "a-b.c d@e!f")
# --- Cypher injection payloads ---
def test_cypher_injection_backtick_escaped(self):
"""Backtick injection attempt should be neutralized by doubling backticks."""
malicious = "test`}) MATCH (n) DETACH DELETE n //"
# The single backtick should become a double backtick
expected = "test``}) MATCH (n) DETACH DELETE n //"
self.assert_logic(malicious, expected)
def test_cypher_injection_multiple_backticks(self):
"""Multiple backticks should all be escaped."""
malicious = "`DROP`DATABASE`"
expected = "``DROP``DATABASE``"
self.assert_logic(malicious, expected)
def test_cypher_injection_curly_braces_preserved(self):
"""Curly brace injection is harmless when enclosed in backticks, so preserved."""
malicious = "test}) RETURN 1 //"
self.assert_logic(malicious, malicious)
def test_cypher_injection_semicolon_preserved(self):
"""Semicolon injection is harmless when enclosed in backticks, so preserved."""
malicious = "test; DROP DATABASE neo4j"
self.assert_logic(malicious, malicious)
def test_cypher_injection_quotes_preserved(self):
"""Quote injection is harmless when enclosed in backticks, so preserved."""
malicious = 'test" OR 1=1 //'
self.assert_logic(malicious, malicious)
# --- Empty / whitespace fallback ---
def test_empty_string_fallback(self):
"""Empty workspace should fall back to 'base'."""
self.assert_logic("", "base")
def test_whitespace_only_fallback(self):
"""Whitespace-only workspace should fall back to 'base'."""
self.assert_logic(" ", "base")
def test_special_chars_only_preserved(self):
"""Workspace with only special characters should be preserved."""
self.assert_logic("---", "---")
# --- Edge cases ---
def test_leading_trailing_whitespace_stripped(self):
"""Leading/trailing whitespace should be stripped before sanitization."""
self.assert_logic(" myworkspace ", "myworkspace")
def test_unicode_characters_preserved(self):
"""Non-ASCII/Chinese characters should be preserved."""
self.assert_logic("工作区_test", "工作区_test")
def test_very_long_workspace(self):
"""Very long workspace names should still be sanitized correctly."""
long_name = "a" * 1000 + "`"
expected = "a" * 1000 + "``"
self.assert_logic(long_name, expected)
def test_single_underscore(self):
"""Single underscore should be valid."""
self.assert_logic("_", "_")
def test_result_always_escapes_backticks(self):
"""Parametric check: any output must not contain unescaped single backticks."""
dangerous_inputs = [
"normal",
"with spaces",
"with-dashes",
"with.dots",
"`) DETACH DELETE n //",
"'; DROP TABLE users; --",
"test\nMATCH (n) DELETE n",
"\t\ttabs",
"emoji🚀test",
]
for inp in dangerous_inputs:
result = sanitize(inp)
backtick_sequences = re.findall(r"`+", result)
for seq in backtick_sequences:
# Any sequence of backticks should have an EVEN length because each ` becomes ``
assert len(seq) % 2 == 0, (
f"Unescaped backtick found in result '{result}' for input '{inp}'"
)