"""
Unit tests for the Treasure Data API client.
"""
import pytest
import requests
import responses
from td_mcp_server.api import Database, Project, Table, TreasureDataClient
class TestTreasureDataClient:
"""Tests for the TreasureDataClient class."""
def setup_method(self):
"""Set up test environment before each test."""
self.api_key = "test_api_key"
self.endpoint = "api.treasuredata.com"
self.client = TreasureDataClient(api_key=self.api_key, endpoint=self.endpoint)
self.mock_databases = [
{
"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,
},
{
"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,
},
{
"name": "db3",
"created_at": "2023-01-03 00:00:00 UTC",
"updated_at": "2023-01-03 00:00:00 UTC",
"count": 0,
"organization": None,
"permission": "administrator",
"delete_protected": False,
},
]
self.mock_tables = [
{
"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,
},
{
"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 = [
{
"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": [],
},
{
"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": [
{"key": "pbp", "value": "cdp_audience"},
{"key": "pbp", "value": "cdp_audience_123456"},
{"key": "sys", "value": "cdp_audience"},
],
},
]
def test_init(self):
"""Test client initialization."""
assert self.client.api_key == self.api_key
assert self.client.endpoint == self.endpoint
assert self.client.base_url == f"https://{self.endpoint}/v3"
assert self.client.headers["Authorization"] == f"TD1 {self.api_key}"
assert self.client.headers["Content-Type"] == "application/json"
def test_init_from_env(self, monkeypatch):
"""Test client initialization from environment variable."""
monkeypatch.setenv("TD_API_KEY", "env_api_key")
client = TreasureDataClient()
assert client.api_key == "env_api_key"
assert client.endpoint == "api.treasuredata.com" # default endpoint
def test_init_missing_api_key(self, monkeypatch):
"""Test client initialization with missing API key."""
monkeypatch.delenv("TD_API_KEY", raising=False)
with pytest.raises(ValueError, match="API key must be provided"):
TreasureDataClient()
@responses.activate
def test_get_databases(self):
"""Test get_databases method."""
# Mock the API response
responses.add(
responses.GET,
f"https://{self.endpoint}/v3/database/list",
json={"databases": self.mock_databases},
status=200,
)
# Call the method
databases = self.client.get_databases()
# Verify the results
assert len(databases) == 3
assert isinstance(databases[0], Database)
assert databases[0].name == "db1"
assert databases[1].name == "db2"
assert databases[2].name == "db3"
assert databases[0].count == 3
assert databases[1].delete_protected is True
@responses.activate
def test_get_databases_with_pagination(self):
"""Test get_databases method with pagination."""
# Mock the API response
responses.add(
responses.GET,
f"https://{self.endpoint}/v3/database/list",
json={"databases": self.mock_databases},
status=200,
)
# Test with limit and offset
databases = self.client.get_databases(limit=2, offset=1)
assert len(databases) == 2
assert databases[0].name == "db2"
assert databases[1].name == "db3"
# Test with all_results=True
databases = self.client.get_databases(all_results=True)
assert len(databases) == 3
# Test with small limit
databases = self.client.get_databases(limit=1)
assert len(databases) == 1
assert databases[0].name == "db1"
@responses.activate
def test_get_database(self):
"""Test get_database method."""
# Mock the API response
responses.add(
responses.GET,
f"https://{self.endpoint}/v3/database/list",
json={"databases": self.mock_databases},
status=200,
)
# Test existing database
database = self.client.get_database("db2")
assert database is not None
assert database.name == "db2"
assert database.count == 5
assert database.delete_protected is True
# Test non-existing database
database = self.client.get_database("nonexistent")
assert database is None
@responses.activate
def test_get_tables(self):
"""Test get_tables method."""
database_name = "test_db"
# Mock the API response
responses.add(
responses.GET,
f"https://{self.endpoint}/v3/table/list/{database_name}",
json={"tables": self.mock_tables},
status=200,
)
# Call the method
tables = self.client.get_tables(database_name)
# Verify the results
assert len(tables) == 2
assert isinstance(tables[0], Table)
assert tables[0].name == "table1"
assert tables[1].name == "table2"
assert tables[0].count == 100
assert tables[1].expire_days == 30
@responses.activate
def test_get_tables_with_pagination(self):
"""Test get_tables method with pagination."""
database_name = "test_db"
# Mock the API response
responses.add(
responses.GET,
f"https://{self.endpoint}/v3/table/list/{database_name}",
json={"tables": self.mock_tables},
status=200,
)
# Test with limit and offset
tables = self.client.get_tables(database_name, limit=1, offset=1)
assert len(tables) == 1
assert tables[0].name == "table2"
# Test with all_results=True
tables = self.client.get_tables(database_name, all_results=True)
assert len(tables) == 2
# Test with large limit
tables = self.client.get_tables(database_name, limit=10)
assert len(tables) == 2
@responses.activate
def test_make_request_error(self):
"""Test error handling in _make_request method."""
# Mock an error response
responses.add(
responses.GET,
f"https://{self.endpoint}/v3/error",
json={"error": "Something went wrong"},
status=500,
)
# Verify that exception is raised
with pytest.raises(requests.exceptions.HTTPError):
self.client._make_request("GET", "error")
@responses.activate
def test_get_projects(self):
"""Test get_projects method."""
# Mock the API response
workflow_endpoint = "api-workflow.treasuredata.com"
responses.add(
responses.GET,
f"https://{workflow_endpoint}/api/projects",
json={"projects": self.mock_projects},
status=200,
)
# Call the method
projects = self.client.get_projects()
# Verify the results
assert len(projects) == 2
assert isinstance(projects[0], Project)
assert projects[0].id == "123456"
assert projects[0].name == "demo_content_affinity"
assert projects[1].id == "789012"
assert projects[1].name == "cdp_audience_123456"
assert len(projects[1].metadata) == 3
assert projects[1].metadata[0].key == "pbp"
assert projects[1].metadata[0].value == "cdp_audience"
@responses.activate
def test_get_projects_with_pagination(self):
"""Test get_projects method with pagination."""
# Mock the API response
workflow_endpoint = "api-workflow.treasuredata.com"
responses.add(
responses.GET,
f"https://{workflow_endpoint}/api/projects",
json={"projects": self.mock_projects},
status=200,
)
# Test with limit and offset
projects = self.client.get_projects(limit=1, offset=1)
assert len(projects) == 1
assert projects[0].id == "789012"
assert projects[0].name == "cdp_audience_123456"
# Test with all_results=True
projects = self.client.get_projects(all_results=True)
assert len(projects) == 2
# Test with large limit
projects = self.client.get_projects(limit=10)
assert len(projects) == 2
def test_workflow_endpoint_derivation(self):
"""Test workflow endpoint derivation based on API endpoint."""
# Test US region standard pattern
client = TreasureDataClient(
api_key=self.api_key, endpoint="api.treasuredata.com"
)
assert client.workflow_endpoint == "api-workflow.treasuredata.com"
# Test Japan region pattern
client = TreasureDataClient(
api_key=self.api_key, endpoint="api.treasuredata.co.jp"
)
assert client.workflow_endpoint == "api-workflow.treasuredata.co.jp"
# Test with non-standard region
client = TreasureDataClient(
api_key=self.api_key, endpoint="api.treasuredata.eu"
)
assert client.workflow_endpoint == "api-workflow.treasuredata.eu"
# Test with different domain structure (non-standard input)
client = TreasureDataClient(
api_key=self.api_key, endpoint="treasuredata-api.com"
)
# Should still perform the replacement
assert client.workflow_endpoint == "treasuredata-api-workflow.com"
# Test custom endpoint
custom_endpoint = "custom-workflow.example.com"
client = TreasureDataClient(
api_key=self.api_key,
endpoint="api.example.com",
workflow_endpoint=custom_endpoint,
)
assert client.workflow_endpoint == custom_endpoint
@responses.activate
def test_get_project(self):
"""Test get_project method."""
project_id = "123456"
workflow_endpoint = "api-workflow.treasuredata.com"
# Mock the API response
responses.add(
responses.GET,
f"https://{workflow_endpoint}/api/projects/{project_id}",
json={
"id": project_id,
"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": [
{"key": "category", "value": "machine-learning"},
{"key": "version", "value": "1.0.0"},
],
},
status=200,
)
# Call the method
project = self.client.get_project(project_id)
# Verify the results
assert project is not None
assert project.id == project_id
assert project.name == "demo_content_affinity"
assert project.revision == "abcdef1234567890abcdef1234567890"
assert project.created_at == "2022-01-01T00:00:00Z"
assert project.updated_at == "2022-01-02T00:00:00Z"
assert len(project.metadata) == 2
assert project.metadata[0].key == "category"
assert project.metadata[0].value == "machine-learning"
@responses.activate
def test_get_project_not_found(self):
"""Test get_project method when project is not found."""
project_id = "nonexistent"
workflow_endpoint = "api-workflow.treasuredata.com"
# Mock the API response with 404 status code
responses.add(
responses.GET,
f"https://{workflow_endpoint}/api/projects/{project_id}",
json={"error": "Project not found"},
status=404,
)
# Call the method - should return None for 404
project = self.client.get_project(project_id)
# Verify the result
assert project is None
@responses.activate
def test_download_project_archive(self, tmp_path):
"""Test download_project_archive method."""
project_id = "123456"
workflow_endpoint = "api-workflow.treasuredata.com"
# Create a temporary file to save the archive
output_path = tmp_path / f"project_{project_id}.tar.gz"
# Mock archive data - creating a simple tar.gz file
mock_archive_data = b"mock tar.gz content"
# Mock the API response
responses.add(
responses.GET,
f"https://{workflow_endpoint}/api/projects/{project_id}/archive",
body=mock_archive_data,
status=200,
)
# Call the method
success = self.client.download_project_archive(project_id, str(output_path))
# Verify the result
assert success is True
assert output_path.exists()
# Check content of the downloaded file
with open(output_path, "rb") as f:
content = f.read()
assert content == mock_archive_data
@responses.activate
def test_download_project_archive_not_found(self, tmp_path):
"""Test download_project_archive method when project is not found."""
project_id = "nonexistent"
workflow_endpoint = "api-workflow.treasuredata.com"
# Create a temporary file to save the archive
output_path = tmp_path / f"project_{project_id}.tar.gz"
# Mock the API response with 404 status code
responses.add(
responses.GET,
f"https://{workflow_endpoint}/api/projects/{project_id}/archive",
json={"error": "Project not found"},
status=404,
)
# Call the method - should return False for 404
success = self.client.download_project_archive(project_id, str(output_path))
# Verify the result
assert success is False
assert not output_path.exists()
@responses.activate
def test_get_workflow_by_id(self):
"""Test get_workflow_by_id method."""
workflow_id = "12345678"
workflow_endpoint = "api-workflow.treasuredata.com"
# Mock the API response
mock_response = {
"id": workflow_id,
"name": "test_workflow",
"project": {
"id": "123456",
"name": "test_project",
},
"revision": "abcdef1234567890abcdef1234567890",
"timezone": "UTC",
"config": {
"+task1": {
"td>": {
"database": "test_db",
"engine": "presto",
}
}
},
}
responses.add(
responses.GET,
f"https://{workflow_endpoint}/api/workflows/{workflow_id}",
json=mock_response,
status=200,
)
# Call the method
workflow = self.client.get_workflow_by_id(workflow_id)
# Verify the result
assert workflow is not None
assert workflow.id == workflow_id
assert workflow.name == "test_workflow"
assert workflow.project.id == "123456"
assert workflow.project.name == "test_project"
assert workflow.timezone == "UTC"
assert workflow.revision == "abcdef1234567890abcdef1234567890"
@responses.activate
def test_get_workflow_by_id_not_found(self):
"""Test get_workflow_by_id method when workflow is not found."""
workflow_id = "nonexistent"
workflow_endpoint = "api-workflow.treasuredata.com"
# Mock the API response with 404 status code
responses.add(
responses.GET,
f"https://{workflow_endpoint}/api/workflows/{workflow_id}",
json={"message": "Resource does not exist: workflow id=nonexistent"},
status=404,
)
# Call the method - should return None for 404
workflow = self.client.get_workflow_by_id(workflow_id)
# Verify the result
assert workflow is None
@responses.activate
def test_get_workflows_with_pagination(self):
"""Test get_workflows method with pagination."""
workflow_endpoint = "api-workflow.treasuredata.com"
# Mock workflow data for page 1
mock_workflows_page1 = [
{
"id": "123",
"name": "workflow1",
"project": {
"id": "1",
"name": "project1",
"updatedAt": "2023-01-01T00:00:00Z",
},
"revision": "abc123",
"timezone": "UTC",
"config": {},
"schedule": None,
"latestSessions": [],
},
{
"id": "456",
"name": "workflow2",
"project": {
"id": "2",
"name": "project2",
"updatedAt": "2023-01-02T00:00:00Z",
},
"revision": "def456",
"timezone": "Asia/Tokyo",
"config": {},
"schedule": {"cron": "0 0 * * *"},
"latestSessions": [],
},
]
# Mock the API response for page 1
responses.add(
responses.GET,
f"https://{workflow_endpoint}/api/console/workflows",
json={"workflows": mock_workflows_page1},
status=200,
match=[
responses.matchers.query_param_matcher(
{
"count": "2",
"page": "1",
"order": "asc",
"sessions": "5",
"output": "simple",
"project_type": "user",
}
)
],
)
# Call the method with specific page
workflows = self.client.get_workflows(count=2, page=1)
# Verify the result
assert len(workflows) == 2
assert workflows[0].id == "123"
assert workflows[0].name == "workflow1"
assert workflows[1].id == "456"
assert workflows[1].name == "workflow2"
@responses.activate
def test_get_session(self):
"""Test get_session method."""
session_id = "123456789"
workflow_endpoint = "api-workflow.treasuredata.com"
# Mock the API response
mock_response = {
"id": session_id,
"project": {"id": "123456", "name": "test_project"},
"workflow": {"name": "test_workflow", "id": "12345678"},
"sessionUuid": "a1b2c3d4-e5f6-7890-abcd-ef1234567890",
"sessionTime": "2025-08-03T03:00:00+00:00",
"lastAttempt": {
"id": "987654321",
"retryAttemptName": None,
"done": True,
"success": True,
"cancelRequested": False,
"params": {},
"createdAt": "2025-08-03T03:00:00Z",
"finishedAt": "2025-08-03T03:05:30Z",
"status": "success",
},
}
responses.add(
responses.GET,
f"https://{workflow_endpoint}/api/sessions/{session_id}",
json=mock_response,
status=200,
)
# Call the method
session = self.client.get_session(session_id)
# Verify the result
assert session is not None
assert session.id == session_id
assert session.project["name"] == "test_project"
assert session.workflow["name"] == "test_workflow"
assert session.session_time == "2025-08-03T03:00:00+00:00"
assert session.last_attempt.success is True
@responses.activate
def test_get_sessions(self):
"""Test get_sessions method."""
workflow_endpoint = "api-workflow.treasuredata.com"
# Mock the API response
mock_response = {
"sessions": [
{
"id": "123456789",
"project": {"id": "123456", "name": "test_project"},
"workflow": {"name": "test_workflow", "id": "12345678"},
"sessionUuid": "a1b2c3d4-e5f6-7890-abcd-ef1234567890",
"sessionTime": "2025-08-03T03:00:00+00:00",
"lastAttempt": {
"id": "987654321",
"retryAttemptName": None,
"done": True,
"success": True,
"cancelRequested": False,
"params": {},
"createdAt": "2025-08-03T03:00:00Z",
"finishedAt": "2025-08-03T03:05:30Z",
"status": "success",
},
}
]
}
responses.add(
responses.GET,
f"https://{workflow_endpoint}/api/sessions",
json=mock_response,
status=200,
match=[responses.matchers.query_param_matcher({"last": "20"})],
)
# Call the method
sessions = self.client.get_sessions()
# Verify the result
assert len(sessions) == 1
assert sessions[0].id == "123456789"
assert sessions[0].project["name"] == "test_project"
@responses.activate
def test_get_attempt(self):
"""Test get_attempt method."""
attempt_id = "987654321"
workflow_endpoint = "api-workflow.treasuredata.com"
# Mock the API response
mock_response = {
"id": attempt_id,
"index": 1,
"project": {"id": "123456", "name": "test_project"},
"workflow": {"name": "test_workflow", "id": "12345678"},
"sessionId": "123456789",
"sessionUuid": "a1b2c3d4-e5f6-7890-abcd-ef1234567890",
"sessionTime": "2025-08-03T03:00:00+00:00",
"retryAttemptName": None,
"done": True,
"success": True,
"cancelRequested": False,
"params": {},
"createdAt": "2025-08-03T03:00:00Z",
"finishedAt": "2025-08-03T03:05:30Z",
"status": "success",
}
responses.add(
responses.GET,
f"https://{workflow_endpoint}/api/attempts/{attempt_id}",
json=mock_response,
status=200,
)
# Call the method
attempt = self.client.get_attempt(attempt_id)
# Verify the result
assert attempt is not None
assert attempt.id == attempt_id
assert attempt.index == 1
assert attempt.session_id == "123456789"
assert attempt.success is True
@responses.activate
def test_get_attempt_tasks(self):
"""Test get_attempt_tasks method."""
attempt_id = "987654321"
workflow_endpoint = "api-workflow.treasuredata.com"
# Mock the API response
mock_response = {
"tasks": [
{
"id": "1234567890",
"fullName": "+main_workflow",
"parentId": None,
"config": {},
"upstreams": [],
"state": "success",
"cancelRequested": False,
"exportParams": {},
"storeParams": {},
"stateParams": {},
"updatedAt": "2025-08-03T03:05:30Z",
"retryAt": None,
"startedAt": "2025-08-03T03:00:00Z",
"error": {},
"isGroup": True,
},
{
"id": "1234567891",
"fullName": "+main_workflow+extract_data",
"parentId": "1234567890",
"config": {
"td>": {
"query": "SELECT * FROM test_table",
"database": "test_db",
}
},
"upstreams": [],
"state": "success",
"cancelRequested": False,
"exportParams": {},
"storeParams": {},
"stateParams": {},
"updatedAt": "2025-08-03T03:02:15Z",
"retryAt": None,
"startedAt": "2025-08-03T03:00:01Z",
"error": {},
"isGroup": False,
},
]
}
responses.add(
responses.GET,
f"https://{workflow_endpoint}/api/attempts/{attempt_id}/tasks",
json=mock_response,
status=200,
)
# Call the method
tasks = self.client.get_attempt_tasks(attempt_id)
# Verify the result
assert len(tasks) == 2
assert tasks[0].id == "1234567890"
assert tasks[0].full_name == "+main_workflow"
assert tasks[0].is_group is True
assert tasks[1].id == "1234567891"
assert tasks[1].full_name == "+main_workflow+extract_data"
assert tasks[1].parent_id == "1234567890"
