"""
Unit tests for the MCP implementation.
"""
import os
from unittest.mock import MagicMock, patch
import pytest
from td_mcp_server.api import Database, Metadata, Project, Table
from td_mcp_server.mcp_impl import (
td_download_project_archive,
td_get_database,
td_get_project,
td_list_databases,
td_list_project_files,
td_list_projects,
td_list_tables,
td_read_project_file,
)
class TestMCPImplementation:
"""Tests for the MCP implementation functions."""
def setup_method(self):
"""Set up test environment before each test."""
self.mock_databases = [
Database(
name="db1",
created_at="2023-01-01 00:00:00 UTC",
updated_at="2023-01-01 00:00:00 UTC",
count=3,
organization=None,
permission="administrator",
delete_protected=False,
),
Database(
name="db2",
created_at="2023-01-02 00:00:00 UTC",
updated_at="2023-01-02 00:00:00 UTC",
count=5,
organization=None,
permission="administrator",
delete_protected=True,
),
]
self.mock_tables = [
Table(
id=1234,
name="table1",
estimated_storage_size=10000,
counter_updated_at="2023-01-01T00:00:00Z",
last_log_timestamp="2023-01-01T00:00:00Z",
delete_protected=False,
created_at="2023-01-01 00:00:00 UTC",
updated_at="2023-01-01 00:00:00 UTC",
type="log",
include_v=True,
count=100,
schema='[["id","string"],["name","string"]]',
expire_days=None,
),
Table(
id=5678,
name="table2",
estimated_storage_size=20000,
counter_updated_at="2023-01-02T00:00:00Z",
last_log_timestamp="2023-01-02T00:00:00Z",
delete_protected=True,
created_at="2023-01-02 00:00:00 UTC",
updated_at="2023-01-02 00:00:00 UTC",
type="log",
include_v=True,
count=200,
schema='[["id","string"],["value","integer"]]',
expire_days=30,
),
]
self.mock_projects = [
Project(
id="123456",
name="demo_content_affinity",
revision="abcdef1234567890abcdef1234567890",
createdAt="2022-01-01T00:00:00Z",
updatedAt="2022-01-02T00:00:00Z",
deletedAt=None,
archiveType="s3",
archiveMd5="abcdefghijklmnopqrstuvwx==",
metadata=[],
),
Project(
id="789012",
name="cdp_audience_123456",
revision="abcdef1234567890abcdef1234567890",
createdAt="2022-01-01T00:00:00Z",
updatedAt="2023-01-01T00:00:00Z",
deletedAt=None,
archiveType="s3",
archiveMd5="zyxwvutsrqponmlkjihgfed==",
metadata=[
Metadata(key="pbp", value="cdp_audience"),
Metadata(key="pbp", value="cdp_audience_123456"),
Metadata(key="sys", value="cdp_audience"),
],
),
]
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_databases_default(self, mock_client_class):
"""Test td_list_databases with default parameters."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_databases.return_value = self.mock_databases
# Call the MCP function
result = await td_list_databases()
# Verify the result
assert "databases" in result
assert result["databases"] == ["db1", "db2"]
assert mock_client.get_databases.called
mock_client.get_databases.assert_called_with(
limit=30, offset=0, all_results=False
)
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_databases_verbose(self, mock_client_class):
"""Test td_list_databases with verbose=True."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_databases.return_value = self.mock_databases
# Call the MCP function
result = await td_list_databases(verbose=True)
# Verify the result
assert "databases" in result
assert len(result["databases"]) == 2
assert result["databases"][0]["name"] == "db1"
assert result["databases"][1]["name"] == "db2"
assert result["databases"][0]["count"] == 3
assert result["databases"][1]["count"] == 5
assert mock_client.get_databases.called
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_databases_pagination(self, mock_client_class):
"""Test td_list_databases with pagination parameters."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_databases.return_value = self.mock_databases
# Call the MCP function
await td_list_databases(limit=10, offset=5, all_results=False)
# Verify the function calls
mock_client.get_databases.assert_called_with(
limit=10, offset=5, all_results=False
)
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_databases_all_results(self, mock_client_class):
"""Test td_list_databases with all_results=True."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_databases.return_value = self.mock_databases
# Call the MCP function
await td_list_databases(all_results=True)
# Verify the function calls
mock_client.get_databases.assert_called_with(
limit=30, offset=0, all_results=True
)
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_databases_no_api_key(self, mock_client_class):
"""Test td_list_databases with no API key."""
# Remove the API key from the environment
with patch.dict(
os.environ, {"TD_API_KEY": "", "TD_ENDPOINT": "api.example.com"}
):
# Call the MCP function
result = await td_list_databases()
# Verify the result contains an error message
assert "error" in result
assert "TD_API_KEY environment variable is not set" in result["error"]
assert not mock_client_class.called
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_get_database(self, mock_client_class):
"""Test td_get_database function."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_database.return_value = self.mock_databases[0]
# Call the MCP function
result = await td_get_database(database_name="db1")
# Verify the result
assert "database" in result
assert result["database"]["name"] == "db1"
assert result["database"]["count"] == 3
assert mock_client.get_database.called
mock_client.get_database.assert_called_with("db1")
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_get_database_not_found(self, mock_client_class):
"""Test td_get_database when database is not found."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_database.return_value = None
# Call the MCP function
result = await td_get_database(database_name="nonexistent")
# Verify the result
assert "error" in result
assert "Database 'nonexistent' not found" in result["error"]
assert mock_client.get_database.called
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_tables_default(self, mock_client_class):
"""Test td_list_tables with default parameters."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_database.return_value = self.mock_databases[0]
mock_client.get_tables.return_value = self.mock_tables
# Call the MCP function
result = await td_list_tables(database_name="db1")
# Verify the result
assert "database" in result
assert "tables" in result
assert result["database"] == "db1"
assert result["tables"] == ["table1", "table2"]
assert mock_client.get_tables.called
mock_client.get_tables.assert_called_with(
"db1", limit=30, offset=0, all_results=False
)
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_tables_verbose(self, mock_client_class):
"""Test td_list_tables with verbose=True."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_database.return_value = self.mock_databases[0]
mock_client.get_tables.return_value = self.mock_tables
# Call the MCP function
result = await td_list_tables(database_name="db1", verbose=True)
# Verify the result
assert "database" in result
assert "tables" in result
assert result["database"] == "db1"
assert len(result["tables"]) == 2
assert result["tables"][0]["name"] == "table1"
assert result["tables"][1]["name"] == "table2"
assert result["tables"][0]["count"] == 100
assert result["tables"][1]["count"] == 200
assert mock_client.get_tables.called
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_tables_pagination(self, mock_client_class):
"""Test td_list_tables with pagination parameters."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_database.return_value = self.mock_databases[0]
mock_client.get_tables.return_value = self.mock_tables
# Call the MCP function
await td_list_tables(database_name="db1", limit=10, offset=5, all_results=False)
# Verify the function calls
mock_client.get_tables.assert_called_with(
"db1", limit=10, offset=5, all_results=False
)
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_tables_all_results(self, mock_client_class):
"""Test td_list_tables with all_results=True."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_database.return_value = self.mock_databases[0]
mock_client.get_tables.return_value = self.mock_tables
# Call the MCP function
await td_list_tables(database_name="db1", all_results=True)
# Verify the function calls
mock_client.get_tables.assert_called_with(
"db1", limit=30, offset=0, all_results=True
)
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_tables_database_not_found(self, mock_client_class):
"""Test td_list_tables when database is not found."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_database.return_value = None
# Call the MCP function
result = await td_list_tables(database_name="nonexistent")
# Verify the result
assert "error" in result
assert "Database 'nonexistent' not found" in result["error"]
assert not mock_client.get_tables.called
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_projects_default(self, mock_client_class):
"""Test td_list_projects with default parameters."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_projects.return_value = self.mock_projects
# Call the MCP function
result = await td_list_projects()
# Verify the result
assert "projects" in result
# Only the non-system project should be included (default behavior)
assert len(result["projects"]) == 1
assert result["projects"][0]["id"] == "123456"
assert result["projects"][0]["name"] == "demo_content_affinity"
assert mock_client.get_projects.called
mock_client.get_projects.assert_called_with(
limit=30, offset=0, all_results=False
)
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_projects_verbose(self, mock_client_class):
"""Test td_list_projects with verbose=True."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_projects.return_value = self.mock_projects
# Call the MCP function
result = await td_list_projects(verbose=True)
# Verify the result
assert "projects" in result
# Only the non-system project should be included (default behavior)
assert len(result["projects"]) == 1
# Check detailed fields in the project
assert result["projects"][0]["id"] == "123456"
assert result["projects"][0]["name"] == "demo_content_affinity"
assert result["projects"][0]["revision"] == "abcdef1234567890abcdef1234567890"
assert result["projects"][0]["created_at"] == "2022-01-01T00:00:00Z"
assert result["projects"][0]["metadata"] == []
# The second project (with "sys" metadata) should be excluded
for project in result["projects"]:
assert project["id"] != "789012"
assert mock_client.get_projects.called
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_projects_pagination(self, mock_client_class):
"""Test td_list_projects with pagination parameters."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_projects.return_value = self.mock_projects
# Call the MCP function
await td_list_projects(limit=10, offset=5, all_results=False)
# Verify the function calls
mock_client.get_projects.assert_called_with(
limit=10, offset=5, all_results=False
)
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_projects_all_results(self, mock_client_class):
"""Test td_list_projects with all_results=True."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_projects.return_value = self.mock_projects
# Call the MCP function
await td_list_projects(all_results=True)
# Verify the function calls
mock_client.get_projects.assert_called_with(
limit=30, offset=0, all_results=True
)
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_projects_exclude_system(self, mock_client_class):
"""Test td_list_projects with system project filtering (default behavior)."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_projects.return_value = self.mock_projects
# Call the MCP function (default is include_system=False)
result = await td_list_projects()
# Verify the result
assert "projects" in result
# Only the first project should be included (the one without "sys" metadata)
assert len(result["projects"]) == 1
assert result["projects"][0]["id"] == "123456"
assert result["projects"][0]["name"] == "demo_content_affinity"
# The second project (with "sys" metadata) should be excluded
for project in result["projects"]:
assert project["id"] != "789012"
assert mock_client.get_projects.called
mock_client.get_projects.assert_called_with(
limit=30, offset=0, all_results=False
)
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_list_projects_include_system(self, mock_client_class):
"""Test td_list_projects with include_system=True."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_projects.return_value = self.mock_projects
# Call the MCP function with include_system=True
result = await td_list_projects(include_system=True)
# Verify the result
assert "projects" in result
# Both projects should be included
assert len(result["projects"]) == 2
# Check that both projects are in the result
project_ids = [p["id"] for p in result["projects"]]
assert "123456" in project_ids
assert "789012" in project_ids
assert mock_client.get_projects.called
mock_client.get_projects.assert_called_with(
limit=30, offset=0, all_results=False
)
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(os.environ, {"TD_API_KEY": "", "TD_ENDPOINT": "api.example.com"})
async def test_td_list_projects_no_api_key(self, mock_client_class):
"""Test td_list_projects with no API key."""
# Call the MCP function
result = await td_list_projects()
# Verify the result contains an error message
assert "error" in result
assert "TD_API_KEY environment variable is not set" in result["error"]
assert not mock_client_class.called
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_get_project(self, mock_client_class):
"""Test td_get_project function."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_project.return_value = self.mock_projects[0]
# Call the MCP function
result = await td_get_project(project_id="123456")
# Verify the result
assert "project" in result
assert result["project"]["id"] == "123456"
assert result["project"]["name"] == "demo_content_affinity"
assert mock_client.get_project.called
mock_client.get_project.assert_called_with("123456")
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_get_project_not_found(self, mock_client_class):
"""Test td_get_project when project is not found."""
# Setup the mock
mock_client = mock_client_class.return_value
mock_client.get_project.return_value = None
# Call the MCP function
result = await td_get_project(project_id="nonexistent")
# Verify the result
assert "error" in result
assert "Project with ID 'nonexistent' not found" in result["error"]
assert mock_client.get_project.called
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch("td_mcp_server.mcp_impl.tempfile.mkdtemp")
@patch("td_mcp_server.mcp_impl.os.chmod")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_download_project_archive(
self, mock_chmod, mock_mkdtemp, mock_client_class
):
"""Test td_download_project_archive with successful download."""
# Setup mocks
mock_temp_dir = "/tmp/td_project_123"
mock_mkdtemp.return_value = mock_temp_dir
mock_client = mock_client_class.return_value
mock_client.get_project.return_value = self.mock_projects[0]
mock_client.download_project_archive.return_value = True
# Call the MCP function
result = await td_download_project_archive(project_id="123456")
# Verify the result
assert "success" in result
assert result["success"] is True
assert result["project_id"] == "123456"
assert result["project_name"] == "demo_content_affinity"
assert result["temp_dir"] == mock_temp_dir
expected_path = os.path.join(mock_temp_dir, "project_123456.tar.gz")
assert result["archive_path"] == expected_path
# Verify API client calls
mock_client.get_project.assert_called_with("123456")
output_path = os.path.join(mock_temp_dir, "project_123456.tar.gz")
mock_client.download_project_archive.assert_called_with("123456", output_path)
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch("td_mcp_server.mcp_impl.tempfile.mkdtemp")
@patch("td_mcp_server.mcp_impl.os.chmod")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_download_project_archive_not_found(
self, mock_chmod, mock_mkdtemp, mock_client_class
):
"""Test td_download_project_archive when project is not found."""
# Setup mocks
mock_temp_dir = "/tmp/td_project_456"
mock_mkdtemp.return_value = mock_temp_dir
mock_client = mock_client_class.return_value
mock_client.get_project.return_value = None
# Call the MCP function
result = await td_download_project_archive(project_id="nonexistent")
# Verify the result
assert "error" in result
assert "Project with ID 'nonexistent' not found" in result["error"]
# Verify that download was not attempted
assert not mock_client.download_project_archive.called
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
@patch("td_mcp_server.mcp_impl.tempfile.mkdtemp")
@patch("td_mcp_server.mcp_impl.os.chmod")
@patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
)
async def test_td_download_project_archive_download_failed(
self, mock_chmod, mock_mkdtemp, mock_client_class
):
"""Test td_download_project_archive when download fails."""
# Setup mocks
mock_temp_dir = "/tmp/td_project_789"
mock_mkdtemp.return_value = mock_temp_dir
mock_client = mock_client_class.return_value
mock_client.get_project.return_value = self.mock_projects[0]
mock_client.download_project_archive.return_value = False
# Call the MCP function
result = await td_download_project_archive(project_id="123456")
# Verify the result
assert "error" in result
assert "Failed to download archive for project '123456'" in result["error"]
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.os.path.exists")
@patch("td_mcp_server.mcp_impl.tarfile.open")
async def test_td_list_project_files(self, mock_tarfile_open, mock_path_exists):
"""Test td_list_project_files successfully listing files."""
# Setup mocks
mock_path_exists.return_value = True
# Create mock tarfile members
mock_file1 = MagicMock()
mock_file1.name = "workflow.dig"
mock_file1.isdir.return_value = False
mock_file1.size = 1024
mock_file2 = MagicMock()
mock_file2.name = "queries/daily_count.sql"
mock_file2.isdir.return_value = False
mock_file2.size = 2048
mock_dir = MagicMock()
mock_dir.name = "queries"
mock_dir.isdir.return_value = True
mock_dir.size = 0
mock_python = MagicMock()
mock_python.name = "scripts/process_data.py"
mock_python.isdir.return_value = False
mock_python.size = 3072
# Setup mock tarfile
mock_tar = MagicMock()
mock_tar.getmembers.return_value = [
mock_file1,
mock_file2,
mock_dir,
mock_python,
]
mock_tarfile_open.return_value.__enter__.return_value = mock_tar
# Call the MCP function
result = await td_list_project_files(archive_path="/tmp/project_123456.tar.gz")
# Verify the result
assert result["success"] is True
assert result["archive_path"] == "/tmp/project_123456.tar.gz"
assert result["file_count"] == 4
# Find each file type and verify its attributes
directory_found = False
workflow_found = False
sql_found = False
python_found = False
for file_info in result["files"]:
if file_info["type"] == "directory" and file_info["name"] == "queries":
directory_found = True
elif file_info["name"] == "workflow.dig":
workflow_found = True
assert file_info["extension"] == ".dig"
assert file_info["file_type"] == "Digdag workflow"
elif file_info["name"] == "queries/daily_count.sql":
sql_found = True
assert file_info["extension"] == ".sql"
assert file_info["file_type"] == "SQL query"
elif file_info["name"] == "scripts/process_data.py":
python_found = True
assert file_info["extension"] == ".py"
assert file_info["file_type"] == "Python script"
assert directory_found, "Directory not found in results"
assert workflow_found, "Workflow file not found in results"
assert sql_found, "SQL file not found in results"
assert python_found, "Python file not found in results"
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.os.path.exists")
async def test_td_list_project_files_not_found(self, mock_path_exists):
"""Test td_list_project_files when archive file not found."""
# Setup mock
mock_path_exists.return_value = False
# Call the MCP function
result = await td_list_project_files(archive_path="/tmp/nonexistent.tar.gz")
# Verify the result
assert "error" in result
assert "Archive file not found" in result["error"]
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.os.path.exists")
@patch("td_mcp_server.mcp_impl.tarfile.open")
async def test_td_read_project_file(self, mock_tarfile_open, mock_path_exists):
"""Test td_read_project_file reading a file successfully."""
# Setup mocks
mock_path_exists.return_value = True
# Create mock tarfile member
mock_file = MagicMock()
mock_file.isdir.return_value = False
mock_file.size = 1024
# Setup mock tar file
mock_tar = MagicMock()
mock_tar.getmember.return_value = mock_file
# Setup mock file content
mock_extracted_file = MagicMock()
mock_extracted_file.read.return_value = (
b"SELECT COUNT(*) FROM events WHERE "
b"td_time_range(time, '2023-01-01', '2023-01-31', 'JST')"
)
mock_tar.extractfile.return_value = mock_extracted_file
mock_tarfile_open.return_value.__enter__.return_value = mock_tar
# Call the MCP function
result = await td_read_project_file(
archive_path="/tmp/project_123456.tar.gz",
file_path="queries/monthly_count.sql",
)
# Verify the result
assert result["success"] is True
assert result["file_path"] == "queries/monthly_count.sql"
assert result["extension"] == ".sql"
assert result["size"] == 1024
assert "SELECT COUNT(*) FROM events" in result["content"]
# Verify tar operations
mock_tar.getmember.assert_called_with("queries/monthly_count.sql")
mock_tar.extractfile.assert_called_with(mock_file)
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.os.path.exists")
async def test_td_read_project_file_archive_not_found(self, mock_path_exists):
"""Test td_read_project_file when archive not found."""
# Setup mock
mock_path_exists.return_value = False
# Call the MCP function
result = await td_read_project_file(
archive_path="/tmp/nonexistent.tar.gz", file_path="workflow.dig"
)
# Verify the result
assert "error" in result
assert "Archive file not found" in result["error"]
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.os.path.exists")
@patch("td_mcp_server.mcp_impl.tarfile.open")
async def test_td_read_project_file_not_found(
self, mock_tarfile_open, mock_path_exists
):
"""Test td_read_project_file when file not in archive."""
# Setup mocks
mock_path_exists.return_value = True
# Setup mock tar file
mock_tar = MagicMock()
mock_tar.getmember.side_effect = KeyError("File not found in archive")
mock_tarfile_open.return_value.__enter__.return_value = mock_tar
# Call the MCP function
result = await td_read_project_file(
archive_path="/tmp/project_123456.tar.gz", file_path="nonexistent.sql"
)
# Verify the result
assert "error" in result
assert "File not found in archive" in result["error"]
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.os.path.exists")
@patch("td_mcp_server.mcp_impl.tarfile.open")
async def test_td_read_project_file_is_directory(
self, mock_tarfile_open, mock_path_exists
):
"""Test td_read_project_file when path is a directory."""
# Setup mocks
mock_path_exists.return_value = True
# Create mock tarfile member (directory)
mock_dir = MagicMock()
mock_dir.isdir.return_value = True
mock_dir.name = "queries"
mock_dir.size = 0
# Setup mock tar file
mock_tar = MagicMock()
mock_tar.getmember.return_value = mock_dir
mock_tarfile_open.return_value.__enter__.return_value = mock_tar
# Call the MCP function
result = await td_read_project_file(
archive_path="/tmp/project_123456.tar.gz", file_path="queries"
)
# Verify the result
assert "error" in result
assert "Cannot read directory contents" in result["error"]
