Spaces:
Running
Running
| 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) | |