Newer
Older
TelosDB / src / backend / tests / integration_test.rs
@楽曲作りまくりおじさん 楽曲作りまくりおじさん 6 days ago 2 KB chore: comprehensive project restructuring and cleanup
use app_lib::{db, llama::LlamaClient, mcp, AppState};
use axum::extract::State;
use axum::Json;
use sea_orm::{ConnectionTrait, DatabaseBackend, DatabaseConnection, Statement};
use serde_json::json;
use std::sync::Arc;

#[tokio::test]
async fn test_mcp_integration_flow() {
    // 1. Initialize DB (in-memory)
    let conn = sea_orm::Database::connect("sqlite::memory:").await.unwrap();

    // Create items table manually for the test
    let db_backend = conn.get_database_backend();
    conn.execute(
        db_backend.build(
            sea_orm::sea_query::Table::create()
                .table(app_lib::entities::items::Entity)
                .if_not_exists()
                .col(
                    sea_orm::sea_query::ColumnDef::new(app_lib::entities::items::Column::Id)
                        .integer()
                        .not_null()
                        .auto_increment()
                        .primary_key(),
                )
                .col(
                    sea_orm::sea_query::ColumnDef::new(app_lib::entities::items::Column::Content)
                        .string()
                        .not_null(),
                )
                .col(
                    sea_orm::sea_query::ColumnDef::new(app_lib::entities::items::Column::Path)
                        .string(),
                )
                .col(
                    sea_orm::sea_query::ColumnDef::new(app_lib::entities::items::Column::CreatedAt)
                        .string(),
                )
                .col(
                    sea_orm::sea_query::ColumnDef::new(app_lib::entities::items::Column::UpdatedAt)
                        .string(),
                ),
        ),
    )
    .await
    .expect("Failed to create table");

    // Create vec_items table (virtual) - might fail if vec0 extension is missing,
    // so we handle it gracefully in the test or skip it.
    // In this test, we just test the tool list and basic error handling which don't need vec0.

    let llama = LlamaClient::new(
        "http://localhost:8080".to_string(),
        "m".to_string(),
        "m".to_string(),
    );

    let state = Arc::new(AppState {
        db: conn,
        llama: Arc::new(llama),
    });

    // 2. Test tools/list
    let req = mcp::JsonRpcRequest {
        jsonrpc: "2.0".to_string(),
        method: "tools/list".to_string(),
        params: json!({}),
        id: Some(json!(1)),
    };

    let res = mcp::message_handler(State(state.clone()), Json(req)).await;
    assert!(res.result.is_some());
    let tools = res.result.as_ref().unwrap()["tools"].as_array().unwrap();
    assert!(tools.len() >= 4);

    // 3. Test Unknown Tool
    let req_unknown = mcp::JsonRpcRequest {
        jsonrpc: "2.0".to_string(),
        method: "tools/call".to_string(),
        params: json!({"name": "unknown"}),
        id: Some(json!(2)),
    };
    let res_err = mcp::message_handler(State(state.clone()), Json(req_unknown)).await;
    assert!(res_err.result.as_ref().unwrap()["error"] == "Unknown tool");
}