arcade-mcp/toolkits/google/tests/test_drive.py
Sam Partee b6b4cd0a4c
🏗️ Restructure: Multi-Package Architecture + uv Migration (#412)
### Overview
Major restructuring from monolithic `arcade-ai` package to modular
library architecture with standardized uv-based dependency management.

![arcade-ai Monorepo
(2)](https://github.com/user-attachments/assets/25f102b0-bb87-4a04-9701-d227d05664b1)

### New Package Structure
- **`arcade-tdk`** - Lightweight toolkit development kit (core
decorators, auth)
- **`arcade-core`** - Core execution engine and catalog functionality  
- **`arcade-serve`** - FastAPI/MCP server components
- **`arcade-ai`** - Meta package that includes CLI functionality.
Optionally include evals via the `evals` extra. Optionally include all
packages via the `all` extra.

### Key Benefits
- **Lighter Dependencies**: Toolkits now depend only on `arcade-tdk` (~2
deps) vs full `arcade-ai` (~30+ deps)
- **Faster Builds**: uv provides 10-100x faster dependency resolution
and installation
- **Better Modularity**: Clear separation of concerns, consumers import
only what they need
- **Standard Tooling**: Eliminates custom poetry scripts, uses standard
Python packaging

### Migration Impact
- All 20 toolkits converted from poetry → uv with `arcade-tdk`
dependencies plus `arcade-ai[evals]` and `arcade-serve` dev
dependencies. When developing locally, devs should install toolkits via
`make install-local`.
- Modern Python 3.10+ type hints throughout
- Standardized build system with hatchling backend
- Enhanced Makefile with robust toolkit management commands
- Removed `arcade dev` CLI command
- Reduce the number of files created by `arcade new` and add an option
to not generate a tests and evals folder.

This foundation enables faster development cycles and cleaner dependency
chains for the growing toolkit ecosystem.

### Todo After this PR is merged
- [ ] Post-merge workflow(s) (release & publish containers, etc)
- [ ] Release order plan. @EricGustin suggests releasing in the
following order:
    1. `arcade-core` version 0.1.0
    2. `arcade-serve` version 0.1.0 and `arcade-tdk` version 0.1.0
    3. `arcade-ai` version 2.0.0
4. Patch release for all toolkits (all changes in toolkits are internal
refactors)
- [ ] [Update docs](https://github.com/ArcadeAI/docs/pull/318)

---------

Co-authored-by: Eric Gustin <eric@arcade.dev>
Co-authored-by: Eric Gustin <34000337+EricGustin@users.noreply.github.com>
2025-06-11 16:48:17 -07:00

391 lines
16 KiB
Python

from unittest.mock import AsyncMock, patch
import pytest
from arcade_tdk.errors import ToolExecutionError
from googleapiclient.errors import HttpError
from arcade_google.models import Corpora, DocumentFormat, OrderBy
from arcade_google.tools import (
get_file_tree_structure,
search_and_retrieve_documents,
search_documents,
)
from arcade_google.utils import build_drive_service
@pytest.fixture
def mock_context():
context = AsyncMock()
context.authorization.token = "mock_token" # noqa: S105
return context
@pytest.fixture
def mock_service():
with patch("arcade_google.tools.drive." + build_drive_service.__name__) as mock_build_service:
yield mock_build_service.return_value
@pytest.mark.asyncio
async def test_search_documents_success(mock_context, mock_service):
# Mock the service.files().list().execute() method
mock_service.files.return_value.list.return_value.execute.side_effect = [
{
"files": [
{"id": "file1", "name": "Document 1"},
{"id": "file2", "name": "Document 2"},
],
"nextPageToken": None,
}
]
result = await search_documents(mock_context, limit=2)
assert result["documents_count"] == 2
assert len(result["documents"]) == 2
assert result["documents"][0]["id"] == "file1"
assert result["documents"][1]["id"] == "file2"
@pytest.mark.asyncio
async def test_search_documents_pagination(mock_context, mock_service):
# Simulate multiple pages
mock_service.files.return_value.list.return_value.execute.side_effect = [
{
"files": [{"id": f"file{i}", "name": f"Document {i}"} for i in range(1, 11)],
"nextPageToken": "token1",
},
{
"files": [{"id": f"file{i}", "name": f"Document {i}"} for i in range(11, 21)],
"nextPageToken": None,
},
]
result = await search_documents(mock_context, limit=15)
assert result["documents_count"] == 15
assert len(result["documents"]) == 15
assert result["documents"][0]["id"] == "file1"
assert result["documents"][-1]["id"] == "file15"
@pytest.mark.asyncio
async def test_search_documents_http_error(mock_context, mock_service):
# Simulate HttpError
mock_service.files.return_value.list.return_value.execute.side_effect = HttpError(
resp=AsyncMock(status=403), content=b'{"error": {"message": "Forbidden"}}'
)
with pytest.raises(
ToolExecutionError, match=f"Error in execution of {search_documents.__tool_name__}"
):
await search_documents(mock_context)
@pytest.mark.asyncio
async def test_search_documents_unexpected_error(mock_context, mock_service):
# Simulate unexpected exception
mock_service.files.return_value.list.return_value.execute.side_effect = Exception(
"Unexpected error"
)
with pytest.raises(
ToolExecutionError, match=f"Error in execution of {search_documents.__tool_name__}"
):
await search_documents(mock_context)
@pytest.mark.asyncio
async def test_search_documents_in_organization_domains(mock_context, mock_service):
# Mock the service.files().list().execute() method
mock_service.files.return_value.list.return_value.execute.side_effect = [
{
"files": [
{"id": "file1", "name": "Document 1"},
],
"nextPageToken": None,
}
]
result = await search_documents(
mock_context,
order_by=OrderBy.MODIFIED_TIME_DESC,
include_shared_drives=False,
include_organization_domain_documents=True,
limit=1,
)
assert result["documents_count"] == 1
mock_service.files.return_value.list.assert_called_with(
q="(mimeType = 'application/vnd.google-apps.document' and trashed = false)",
corpora=Corpora.DOMAIN.value,
pageSize=1,
orderBy=OrderBy.MODIFIED_TIME_DESC.value,
includeItemsFromAllDrives="true",
supportsAllDrives="true",
)
@pytest.mark.asyncio
@patch("arcade_google.tools.drive.search_documents")
@patch("arcade_google.tools.drive.get_document_by_id")
async def test_search_and_retrieve_documents_in_markdown_format(
mock_get_document_by_id,
mock_search_documents,
mock_context,
sample_document_and_expected_formats,
):
(sample_document, expected_markdown, _) = sample_document_and_expected_formats
mock_search_documents.return_value = {
"documents_count": 1,
"documents": [{"id": sample_document["documentId"], "title": sample_document["title"]}],
}
mock_get_document_by_id.return_value = sample_document
result = await search_and_retrieve_documents(
mock_context,
document_contains=[sample_document["title"]],
return_format=DocumentFormat.MARKDOWN,
)
assert result["documents_count"] == 1
assert result["documents"][0] == expected_markdown
@pytest.mark.asyncio
@patch("arcade_google.tools.drive.search_documents")
@patch("arcade_google.tools.drive.get_document_by_id")
async def test_search_and_retrieve_documents_in_html_format(
mock_get_document_by_id,
mock_search_documents,
mock_context,
sample_document_and_expected_formats,
):
(sample_document, _, expected_html) = sample_document_and_expected_formats
mock_search_documents.return_value = {
"documents_count": 1,
"documents": [{"id": sample_document["documentId"], "title": sample_document["title"]}],
}
mock_get_document_by_id.return_value = sample_document
result = await search_and_retrieve_documents(
mock_context,
document_contains=[sample_document["title"]],
return_format=DocumentFormat.HTML,
)
assert result["documents_count"] == 1
assert result["documents"][0] == expected_html
@pytest.mark.asyncio
@patch("arcade_google.tools.drive.search_documents")
@patch("arcade_google.tools.drive.get_document_by_id")
async def test_search_and_retrieve_documents_in_google_json_format(
mock_get_document_by_id,
mock_search_documents,
mock_context,
sample_document_and_expected_formats,
):
(sample_document, _, _) = sample_document_and_expected_formats
mock_search_documents.return_value = {
"documents_count": 1,
"documents": [{"id": sample_document["documentId"], "title": sample_document["title"]}],
}
mock_get_document_by_id.return_value = sample_document
result = await search_and_retrieve_documents(
mock_context,
document_contains=[sample_document["title"]],
return_format=DocumentFormat.GOOGLE_API_JSON,
)
assert result["documents_count"] == 1
assert result["documents"][0] == sample_document
@pytest.mark.asyncio
async def test_get_file_tree_structure(
mock_context, mock_service, sample_drive_file_tree_request_responses
):
files_list_sample, drives_get_sample = sample_drive_file_tree_request_responses
mock_service.files.return_value.list.return_value.execute.side_effect = [files_list_sample]
mock_service.drives.return_value.get.return_value.execute.side_effect = drives_get_sample
result = await get_file_tree_structure(mock_context, include_shared_drives=True)
expected_file_tree = {
"drives": [
{
"id": "0AFqcR6obkydtUk9PVA",
"name": "Shared Drive 1",
"children": [
{
"createdTime": "2025-02-26T00:27:45.526Z",
"id": "1dCOCdPxhTqiB3j3bWrIWM692ZbL8dyjt",
"mimeType": "application/vnd.google-apps.folder",
"modifiedTime": "2025-02-26T00:27:45.526Z",
"name": "shared-1-folder-1",
"children": [
{
"createdTime": "2025-02-26T00:28:20.571Z",
"id": "19WVyQndQsc0AxxfdrIt5CvDQd6r-BvpqnB8bWZoL7Xk",
"mimeType": "application/vnd.google-apps.document",
"modifiedTime": "2025-02-26T00:28:30.773Z",
"name": "shared-1-folder-1-doc-1",
"size": {
"unit": "bytes",
"value": 1024,
},
}
],
},
{
"createdTime": "2025-02-26T00:27:19.287Z",
"id": "1didt_h-tDjuJ-dmYtHUSyOCPci30K_kSszvg0G3tKBM",
"mimeType": "application/vnd.google-apps.document",
"modifiedTime": "2025-02-26T00:27:26.079Z",
"name": "shared-1-doc-1",
"size": {
"unit": "bytes",
"value": 1024,
},
},
],
},
{
"name": "My Drive",
"children": [
{
"createdTime": "2025-01-24T06:34:22.305Z",
"id": "1vB6sv0MD0hYSraYvWU_fcci3GN_-Jf4g-LfyXdG8ZMo",
"mimeType": "application/vnd.google-apps.document",
"modifiedTime": "2025-02-25T21:54:30.632Z",
"name": "The Birth of MX Engineering",
"owners": [
{
"email": "one_new_tool_everyday@arcade.dev",
"name": "one_new_tool_everyday",
}
],
"size": {
"unit": "bytes",
"value": 6634,
},
},
{
"createdTime": "2025-02-25T17:57:46.036Z",
"id": "1gqioaHG53jPVeJN5gBpHoO-GWtwiJcLo",
"mimeType": "application/vnd.google-apps.folder",
"modifiedTime": "2025-02-25T17:57:46.036Z",
"name": "test folder 1",
"owners": [
{
"email": "one_new_tool_everyday@arcade.dev",
"name": "one_new_tool_everyday",
}
],
"children": [
{
"id": "1J92V9yvVWm_uNHq3CCY4wyG1H9B6iiwO",
"name": "test folder 1.1",
"mimeType": "application/vnd.google-apps.folder",
"createdTime": "2025-02-25T17:58:58.987Z",
"modifiedTime": "2025-02-25T17:58:58.987Z",
"owners": [
{
"email": "one_new_tool_everyday@arcade.dev",
"name": "one_new_tool_everyday",
}
],
"children": [
{
"id": "1wv2dmYo0skJTI59ZIcwH9vm-wt7psMwXTvihuEGeHeI",
"name": "test document 1.1.1",
"mimeType": "application/vnd.google-apps.document",
"createdTime": "2025-02-25T17:59:03.325Z",
"modifiedTime": "2025-02-25T17:59:11.445Z",
"owners": [
{
"email": "one_new_tool_everyday@arcade.dev",
"name": "one_new_tool_everyday",
}
],
"size": {
"unit": "bytes",
"value": 1024,
},
},
],
},
{
"id": "1DSmL7d07kjT6b6L-t4JIT06ElUbZ1q0K6_gEpn_UGZ8",
"name": "test document 1.2",
"mimeType": "application/vnd.google-apps.document",
"createdTime": "2025-02-25T17:58:38.628Z",
"modifiedTime": "2025-02-25T17:58:46.713Z",
"owners": [
{
"email": "one_new_tool_everyday@arcade.dev",
"name": "one_new_tool_everyday",
}
],
"size": {
"unit": "bytes",
"value": 1024,
},
},
{
"id": "1Fcxz7HsyO2Zyc-5DTD3zBQnaVrZwD29BP9KD9rPnYfE",
"name": "test document 1.1",
"mimeType": "application/vnd.google-apps.document",
"createdTime": "2025-02-25T17:57:53.850Z",
"modifiedTime": "2025-02-25T17:58:28.745Z",
"owners": [
{
"email": "one_new_tool_everyday@arcade.dev",
"name": "one_new_tool_everyday",
}
],
"size": {
"unit": "bytes",
"value": 1024,
},
},
],
},
{
"createdTime": "2025-02-18T20:48:52.786Z",
"id": "16PUe97yGQeOjQgrgd54iCoxzid4SEvu_J33P_ELd5r8",
"mimeType": "application/vnd.google-apps.presentation",
"modifiedTime": "2025-02-19T23:31:20.483Z",
"name": "Hello world presentation",
"owners": [
{
"email": "john.doe@arcade.dev",
"name": "john.doe",
}
],
"size": {
"unit": "bytes",
"value": 15774558,
},
},
{
"id": "1nG7lSvIyK05N9METPczVJa4iGgE7uoo-A6zpqjpUsDY",
"name": "Shared doc 1",
"mimeType": "application/vnd.google-apps.document",
"createdTime": "2025-02-19T18:51:44.622Z",
"modifiedTime": "2025-02-19T19:30:39.773Z",
"owners": [
{
"name": "theboss",
"email": "theboss@arcade.dev",
}
],
"size": {
"unit": "bytes",
"value": 2700,
},
},
],
},
]
}
assert result == expected_file_tree