SwiftDBAI/Tests/SwiftDBAITests/DatabaseToolTests.swift

// DatabaseToolTests.swift
// SwiftDBAI

import Testing
import Foundation
import GRDB
@testable import SwiftDBAI

@Suite("DatabaseTool")
struct DatabaseToolTests {

    // MARK: - Helper

    /// Creates an in-memory database with sample data for testing.
    private func makeTestDatabase() throws -> DatabaseQueue {
        let db = try DatabaseQueue(configuration: {
            var config = Configuration()
            config.foreignKeysEnabled = true
            return config
        }())

        try db.write { db in
            try db.execute(sql: """
                CREATE TABLE users (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL,
                    email TEXT UNIQUE
                );
                """)

            try db.execute(sql: """
                CREATE TABLE posts (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id INTEGER NOT NULL REFERENCES users(id),
                    title TEXT NOT NULL,
                    body TEXT,
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP
                );
                """)

            try db.execute(sql: """
                CREATE INDEX idx_posts_user ON posts(user_id);
                """)

            // Insert sample data
            try db.execute(sql: "INSERT INTO users (name, email) VALUES ('Alice', 'alice@example.com')")
            try db.execute(sql: "INSERT INTO users (name, email) VALUES ('Bob', 'bob@example.com')")
            try db.execute(sql: "INSERT INTO posts (user_id, title, body) VALUES (1, 'Hello World', 'First post')")
            try db.execute(sql: "INSERT INTO posts (user_id, title, body) VALUES (1, 'Second Post', 'More content')")
            try db.execute(sql: "INSERT INTO posts (user_id, title, body) VALUES (2, 'Bob Post', 'Bob writes')")
        }

        return db
    }

    // MARK: - Creation

    @Test("Creates tool from database connection")
    func testCreationFromDatabase() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db)

        #expect(tool.name == "execute_sql")
        #expect(!tool.description.isEmpty)
    }

    @Test("Creates tool from database path")
    func testCreationFromPath() async throws {
        let tmpDir = FileManager.default.temporaryDirectory
        let dbPath = tmpDir.appendingPathComponent("test_\(UUID().uuidString).sqlite").path
        defer { try? FileManager.default.removeItem(atPath: dbPath) }

        // Create a database at the path
        let dbQueue = try DatabaseQueue(path: dbPath)
        try await dbQueue.write { db in
            try db.execute(sql: "CREATE TABLE test (id INTEGER PRIMARY KEY)")
        }

        let tool = try await DatabaseTool(databasePath: dbPath)
        #expect(tool.name == "execute_sql")
    }

    // MARK: - Schema Context

    @Test("Schema context contains table info")
    func testSchemaContext() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db)

        let context = tool.schemaContext
        #expect(context.contains("users"))
        #expect(context.contains("posts"))
        #expect(context.contains("name"))
        #expect(context.contains("email"))
    }

    @Test("System prompt snippet contains schema")
    func testSystemPromptSnippet() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db)

        let snippet = tool.systemPromptSnippet
        #expect(snippet.contains("users"))
        #expect(snippet.contains("posts"))
        #expect(snippet.contains("execute_sql"))
        #expect(snippet.contains("SELECT"))
    }

    // MARK: - SQL Execution

    @Test("Executes valid SELECT query")
    func testExecuteSelect() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db)

        let result = try tool.execute(sql: "SELECT name, email FROM users ORDER BY name")

        #expect(result.rowCount == 2)
        #expect(result.columns == ["name", "email"])
        #expect(result.rows[0]["name"] == "Alice")
        #expect(result.rows[1]["name"] == "Bob")
    }

    @Test("Executes query with JOIN")
    func testExecuteJoin() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db)

        let result = try tool.execute(sql: """
            SELECT u.name, COUNT(p.id) as post_count
            FROM users u
            JOIN posts p ON p.user_id = u.id
            GROUP BY u.name
            ORDER BY u.name
            """)

        #expect(result.rowCount == 2)
        #expect(result.rows[0]["name"] == "Alice")
        #expect(result.rows[0]["post_count"] == "2")
    }

    @Test("Rejects INSERT with read-only allowlist")
    func testRejectInsert() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db, allowlist: .readOnly)

        #expect(throws: SQLParsingError.self) {
            try tool.execute(sql: "INSERT INTO users (name, email) VALUES ('Eve', 'eve@example.com')")
        }
    }

    @Test("Rejects DELETE with read-only allowlist")
    func testRejectDelete() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db, allowlist: .readOnly)

        #expect(throws: SQLParsingError.self) {
            try tool.execute(sql: "DELETE FROM users WHERE id = 1")
        }
    }

    @Test("Rejects DROP as dangerous operation")
    func testRejectDrop() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db, allowlist: .unrestricted)

        #expect(throws: SQLParsingError.self) {
            try tool.execute(sql: "DROP TABLE users")
        }
    }

    @Test("Executes raw query returning QueryResult")
    func testExecuteRaw() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db)

        let result = try tool.executeRaw(sql: "SELECT COUNT(*) as cnt FROM users")

        #expect(result.rowCount == 1)
        #expect(result.columns == ["cnt"])
        if case .integer(let count) = result.rows[0]["cnt"] {
            #expect(count == 2)
        } else {
            Issue.record("Expected integer value")
        }
    }

    // MARK: - ToolResult Formatting

    @Test("ToolResult JSON serialization")
    func testToolResultJSON() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db)

        let result = try tool.execute(sql: "SELECT name FROM users ORDER BY name LIMIT 1")

        let json = result.jsonString
        #expect(json.contains("\"columns\""))
        #expect(json.contains("\"rows\""))
        #expect(json.contains("\"row_count\""))
        #expect(json.contains("Alice"))

        // Verify it is valid JSON
        let data = json.data(using: .utf8)!
        let parsed = try JSONSerialization.jsonObject(with: data) as! [String: Any]
        #expect(parsed["row_count"] as? Int == 1)
    }

    @Test("ToolResult markdown table formatting")
    func testToolResultMarkdown() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db)

        let result = try tool.execute(sql: "SELECT name, email FROM users ORDER BY name")

        let md = result.markdownTable
        #expect(md.contains("| name | email |"))
        #expect(md.contains("| --- | --- |"))
        #expect(md.contains("| Alice | alice@example.com |"))
        #expect(md.contains("| Bob | bob@example.com |"))
    }

    @Test("ToolResult markdown table with empty result")
    func testToolResultMarkdownEmpty() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db)

        let result = try tool.execute(sql: "SELECT name FROM users WHERE name = 'Nobody'")

        #expect(result.markdownTable == "_No results._")
    }

    @Test("ToolResult text summary")
    func testToolResultTextSummary() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db)

        let result = try tool.execute(sql: "SELECT name FROM users")

        let summary = result.textSummary
        #expect(summary.contains("2 rows"))
        #expect(summary.contains("name"))
    }

    @Test("ToolResult text summary with empty result")
    func testToolResultTextSummaryEmpty() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db)

        let result = try tool.execute(sql: "SELECT name FROM users WHERE 1 = 0")

        #expect(result.textSummary.contains("no results"))
    }

    // MARK: - Parameters Schema

    @Test("Parameters schema has correct structure")
    func testParametersSchema() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db)

        let schema = tool.parametersSchema
        #expect(schema["type"] as? String == "object")

        let properties = schema["properties"] as? [String: Any]
        #expect(properties != nil)

        let sqlProp = properties?["sql"] as? [String: Any]
        #expect(sqlProp?["type"] as? String == "string")

        let required = schema["required"] as? [String]
        #expect(required == ["sql"])
    }

    // MARK: - OpenAI Function Definition

    @Test("OpenAI function definition has correct format")
    func testOpenAIFunctionDefinition() async throws {
        let db = try makeTestDatabase()
        let tool = try await DatabaseTool(database: db)

        let def = tool.openAIFunctionDefinition
        #expect(def["type"] as? String == "function")

        let function = def["function"] as? [String: Any]
        #expect(function?["name"] as? String == "execute_sql")
        #expect(function?["description"] as? String != nil)
        #expect(function?["parameters"] as? [String: Any] != nil)

        // Verify it can be serialized to JSON
        let data = try JSONSerialization.data(withJSONObject: def)
        #expect(data.count > 0)
    }

    // MARK: - ToolResult Codable

    @Test("ToolResult is Codable")
    func testToolResultCodable() throws {
        let result = ToolResult(
            columns: ["name", "age"],
            rows: [["name": "Alice", "age": "30"]],
            rowCount: 1,
            executionTime: 0.005,
            sql: "SELECT name, age FROM users"
        )

        let encoder = JSONEncoder()
        let data = try encoder.encode(result)
        let decoder = JSONDecoder()
        let decoded = try decoder.decode(ToolResult.self, from: data)

        #expect(decoded.columns == result.columns)
        #expect(decoded.rows == result.rows)
        #expect(decoded.rowCount == result.rowCount)
        #expect(decoded.sql == result.sql)
    }
}