ollama-proxy 0.1.0

Proxy server for Ollama
//! Test suite for the database logger module.

#[cfg(feature = "database-logging")]
mod tests {
    use crate::db_logger::DbLogger;
    use hyper::{HeaderMap, Method, StatusCode};
    use std::net::{IpAddr, Ipv4Addr, SocketAddr};
    use tempfile::tempdir;
    use tokio::test;

    #[test]
    async fn test_no_database_url() {
        // Create a logger without a database URL
        let logger = DbLogger::new(None).await;

        // Create test data
        let client_ip = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080);
        let method = Method::GET;
        let path = "/api/models";
        let headers = HeaderMap::new();
        let body = vec![];

        // These should not panic even without a database
        logger
            .log_request(&client_ip, &method, path, &headers, &body)
            .await;
        logger
            .log_response(&client_ip, &method, path, &StatusCode::OK, &headers, &body)
            .await;

        // There's no easy way to verify nothing happened, but the test passes if there's no panic
    }

    // The remaining tests we'll manually test in a more complete integration test
    // since they depend on checking actual database contents which requires accessing
    // the internal pool of the DbLogger (which is not exposed)

    // For real integration tests, we'd use a temporary file-based database
    // This is just a placeholder to show how you would do it
    #[test]
    async fn mock_integration_test() {
        // Create a temporary directory for the SQLite database
        let temp_dir = tempdir().unwrap();
        let db_path = temp_dir.path().join("test.db");
        let db_url = format!("sqlite:{}", db_path.display());

        // Initialize the logger with the database URL
        let logger = DbLogger::new(Some(db_url.clone())).await;

        // Log some requests and responses
        let client_ip = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080);
        let method = Method::POST;
        let path = "/api/generate";
        let mut headers = HeaderMap::new();
        headers.insert("content-type", "application/json".parse().unwrap());
        let req_body = r#"{"model": "llama2", "prompt": "Hello world"}"#.as_bytes();

        logger
            .log_request(&client_ip, &method, path, &headers, req_body)
            .await;
        logger
            .log_response(
                &client_ip,
                &method,
                path,
                &StatusCode::OK,
                &headers,
                r#"{"model": "llama2", "response": "Hello human"}"#.as_bytes(),
            )
            .await;

        // In a real test, we'd check the database contents here
        // For this test, we're just ensuring nothing panics
    }
}