Spaces:
Running
Running
File size: 3,818 Bytes
a8c9ee8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | import pytest
from unittest.mock import MagicMock, patch
from backend.SQL_Agent.tenant_file_toolkit import TenantFileToolkit
class TestTenantFileToolkit:
def test_get_tenant_id_fallback(self):
toolkit = TenantFileToolkit(storage_service=MagicMock())
# Since we are not running inside an Agno run context, it should fallback
tenant_id = toolkit._get_tenant_id()
assert tenant_id == "unknown_tenant"
def test_list_tenant_assets_formats_correctly(self):
mock_run_context = MagicMock()
mock_run_context.session_state = {"tenant_id": "tenant_123"}
mock_storage = MagicMock()
mock_storage.client = MagicMock() # bypass truthiness check
# Mock objects returned by minio list_objects
mock_obj1 = MagicMock()
mock_obj1.object_name = "tenant_123/raw/dataset.csv"
mock_obj1.size = 1024 * 1024 * 5 # 5MB
mock_obj1.last_modified.strftime.return_value = "2025-01-01"
mock_obj2 = MagicMock()
mock_obj2.object_name = "tenant_123/analysis/report.json"
mock_obj2.size = 1024 * 500 # 0.5 MB
mock_obj2.last_modified.strftime.return_value = "2025-01-02"
mock_storage.client.list_objects.return_value = [mock_obj1, mock_obj2]
toolkit = TenantFileToolkit(storage_service=mock_storage)
result = toolkit.list_tenant_assets(run_context=mock_run_context)
assert "Assets for Tenant Workspace" in result
assert "Datasets (CSV/Excel/Parquet)" in result
assert "dataset.csv" in result
assert "Reports (.json/.md/.txt/.html)" in result
assert "report.json" in result
def test_list_tenant_assets_surfaces_api_error_when_fallback_empty(self):
mock_run_context = MagicMock()
mock_run_context.session_state = {
"tenant_id": "tenant_123",
"supabase_jwt": "bad_or_expired_token",
}
mock_storage = MagicMock()
mock_storage.client = MagicMock()
mock_storage.client.list_objects.return_value = []
toolkit = TenantFileToolkit(storage_service=mock_storage)
with patch.object(
toolkit,
"_fetch_assets_from_api",
return_value={"error": "Tenant files API request failed (401): unauthorized"},
):
result = toolkit.list_tenant_assets(run_context=mock_run_context)
assert "Unable to reliably list tenant assets right now" in result
assert "401" in result
def test_load_tenant_file_security_check(self):
mock_run_context = MagicMock()
mock_run_context.session_state = {"tenant_id": "tenant_123"}
toolkit = TenantFileToolkit(storage_service=MagicMock())
# Try to read outside the tenant
result = toolkit.load_tenant_file_to_dataframe("tenant_999/files/secret.csv", run_context=mock_run_context)
assert "Access Denied" in result
def test_load_tenant_file_chunking(self):
mock_run_context = MagicMock()
mock_run_context.session_state = {"tenant_id": "tenant_123"}
mock_storage = MagicMock()
toolkit = TenantFileToolkit(storage_service=mock_storage)
# Mock chunked iterator yielding one dataframe
import pandas as pd
df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
mock_storage.load_dataframe.return_value = iter([df])
result = toolkit.load_tenant_file_to_dataframe("tenant_123/files/data.csv", chunksize=100, run_context=mock_run_context)
assert "Successfully read FIRST CHUNK" in result
assert "2 (in chunk)" in result
assert "**Columns:** 2" in result
assert mock_storage.load_dataframe.called_once_with("tenant_123/files/data.csv", chunksize=100)
|