use serde_json::{Value, json};
use task_graph_mcp::config::{IdsConfig, StatesConfig};
use task_graph_mcp::db::Database;
use task_graph_mcp::format::{OutputFormat, ToolResult};
use task_graph_mcp::tools::query;
fn unwrap_json(result: ToolResult) -> Value {
match result {
ToolResult::Json(v) => v,
ToolResult::Raw(s) => panic!("Expected JSON, got raw text: {}", s),
}
}
fn unwrap_raw(result: ToolResult) -> String {
match result {
ToolResult::Raw(s) => s,
ToolResult::Json(v) => panic!("Expected raw text, got JSON: {}", v),
}
}
fn setup_db() -> Database {
Database::open_in_memory().expect("Failed to create in-memory database")
}
fn default_states_config() -> StatesConfig {
StatesConfig::default()
}
#[test]
fn query_select_all_tasks() {
let db = setup_db();
let states_config = default_states_config();
db.create_task(
None,
"Task 1".to_string(),
None,
None,
None, None,
None,
None,
None,
None,
None,
&states_config,
&IdsConfig::default(),
)
.unwrap();
db.create_task(
None,
"Task 2".to_string(),
None,
None,
None, None,
None,
None,
None,
None,
None,
&states_config,
&IdsConfig::default(),
)
.unwrap();
let result = unwrap_json(
query::query(
&db,
OutputFormat::Json,
json!({
"sql": "SELECT id, title FROM tasks ORDER BY created_at"
}),
)
.unwrap(),
);
assert_eq!(result["row_count"], 2);
assert!(!result["truncated"].as_bool().unwrap());
assert_eq!(result["columns"], json!(["id", "title"]));
let rows = result["rows"].as_array().unwrap();
assert_eq!(rows.len(), 2);
assert_eq!(rows[0]["title"], "Task 1");
assert_eq!(rows[1]["title"], "Task 2");
}
#[test]
fn query_with_parameters() {
let db = setup_db();
let states_config = default_states_config();
let task = db
.create_task(
None,
"Find Me".to_string(),
None,
None,
None, None,
None,
None,
None,
None,
None,
&states_config,
&IdsConfig::default(),
)
.unwrap();
db.create_task(
None,
"Other Task".to_string(),
None,
None,
None, None,
None,
None,
None,
None,
None,
&states_config,
&IdsConfig::default(),
)
.unwrap();
let result = unwrap_json(
query::query(
&db,
OutputFormat::Json,
json!({
"sql": "SELECT id, title FROM tasks WHERE id = ?",
"params": [task.id]
}),
)
.unwrap(),
);
assert_eq!(result["row_count"], 1);
let rows = result["rows"].as_array().unwrap();
assert_eq!(rows[0]["title"], "Find Me");
}
#[test]
fn query_enforces_row_limit() {
let db = setup_db();
let states_config = default_states_config();
for i in 0..10 {
db.create_task(
None,
format!("Task {}", i),
None,
None,
None, None,
None,
None,
None,
None,
None,
&states_config,
&IdsConfig::default(),
)
.unwrap();
}
let result = unwrap_json(
query::query(
&db,
OutputFormat::Json,
json!({
"sql": "SELECT id, title FROM tasks",
"limit": 5
}),
)
.unwrap(),
);
assert_eq!(result["row_count"], 5);
assert!(result["truncated"].as_bool().unwrap());
assert_eq!(result["limit"], 5);
}
#[test]
fn query_csv_format() {
let db = setup_db();
let states_config = default_states_config();
db.create_task(
None,
"CSV Task".to_string(),
None,
None,
None, None,
None,
None,
None,
None,
None,
&states_config,
&IdsConfig::default(),
)
.unwrap();
let csv_data = unwrap_raw(
query::query(
&db,
OutputFormat::Json,
json!({
"sql": "SELECT title, status FROM tasks",
"format": "csv"
}),
)
.unwrap(),
);
assert!(csv_data.contains("title,status"));
assert!(csv_data.contains("CSV Task"));
assert!(csv_data.contains("pending")); }
#[test]
fn query_markdown_format() {
let db = setup_db();
let states_config = default_states_config();
db.create_task(
None,
"Markdown Task".to_string(),
None,
None,
None, None,
None,
None,
None,
None,
None,
&states_config,
&IdsConfig::default(),
)
.unwrap();
let md_data = unwrap_raw(
query::query(
&db,
OutputFormat::Json,
json!({
"sql": "SELECT title FROM tasks",
"format": "markdown"
}),
)
.unwrap(),
);
assert!(md_data.contains("| title |"));
assert!(md_data.contains("| --- |"));
assert!(md_data.contains("| Markdown Task |"));
}
#[test]
fn query_rejects_insert() {
let db = setup_db();
let result = query::query(
&db,
OutputFormat::Json,
json!({
"sql": "INSERT INTO tasks (id, title) VALUES ('x', 'bad')"
}),
);
assert!(result.is_err());
let err = result.unwrap_err().to_string();
assert!(err.contains("INSERT") || err.contains("SELECT"));
}
#[test]
fn query_rejects_update() {
let db = setup_db();
let result = query::query(
&db,
OutputFormat::Json,
json!({
"sql": "UPDATE tasks SET title = 'hacked'"
}),
);
assert!(result.is_err());
}
#[test]
fn query_rejects_delete() {
let db = setup_db();
let result = query::query(
&db,
OutputFormat::Json,
json!({
"sql": "DELETE FROM tasks"
}),
);
assert!(result.is_err());
}
#[test]
fn query_rejects_drop() {
let db = setup_db();
let result = query::query(
&db,
OutputFormat::Json,
json!({
"sql": "DROP TABLE tasks"
}),
);
assert!(result.is_err());
}
#[test]
fn query_rejects_multiple_statements() {
let db = setup_db();
let result = query::query(
&db,
OutputFormat::Json,
json!({
"sql": "SELECT 1; DROP TABLE tasks;"
}),
);
assert!(result.is_err());
let err = result.unwrap_err().to_string();
assert!(err.contains("Multiple"));
}
#[test]
fn query_allows_column_names_with_keywords() {
let db = setup_db();
let result = query::query(
&db,
OutputFormat::Json,
json!({
"sql": "SELECT id, status FROM tasks WHERE status = 'pending'"
}),
);
assert!(result.is_ok());
}
#[test]
fn query_with_cte() {
let db = setup_db();
let states_config = default_states_config();
db.create_task(
None,
"CTE Task".to_string(),
None,
None,
None, None,
None,
None,
None,
None,
None,
&states_config,
&IdsConfig::default(),
)
.unwrap();
let result = unwrap_json(
query::query(
&db,
OutputFormat::Json,
json!({
"sql": "WITH task_list AS (SELECT id, title FROM tasks) SELECT * FROM task_list"
}),
)
.unwrap(),
);
assert_eq!(result["row_count"], 1);
}
#[test]
fn query_max_limit_enforced() {
let db = setup_db();
let result = unwrap_json(
query::query(
&db,
OutputFormat::Json,
json!({
"sql": "SELECT 1",
"limit": 5000
}),
)
.unwrap(),
);
assert_eq!(result["limit"], 1000);
}
#[test]
fn query_default_limit() {
let db = setup_db();
let result = unwrap_json(
query::query(
&db,
OutputFormat::Json,
json!({
"sql": "SELECT 1"
}),
)
.unwrap(),
);
assert_eq!(result["limit"], 100);
}