Fundamental_Analysis/services/data-persistence-service/tests/api_tests.rs

#![allow(unused_imports)]

use axum::{
    body::Body,
    http::{self, Request, StatusCode},
    response::Response,
};
use data_persistence_service::{
    self as app,
    dtos::{
        AnalysisResultDto, CompanyProfileDto, DailyMarketDataBatchDto, DailyMarketDataDto,
        NewAnalysisResult, TimeSeriesFinancialBatchDto, TimeSeriesFinancialDto,
    },
    AppState,
};
use http_body_util::BodyExt;
use sqlx::PgPool;
use tower::util::ServiceExt; // for `oneshot`
use uuid::Uuid;

// Note: We need to make `build_rest_router` and `AppState` public in lib.rs and main.rs respectively.
// This test structure assumes that has been done.

#[sqlx::test]
async fn test_api_upsert_and_get_company(pool: PgPool) {
    let state = AppState::new(pool);
    let openapi = app::build_openapi_spec();
    let app = app::build_rest_router_with_state(openapi, state).unwrap();

    // 1. Act: Upsert a new company
    let new_company = CompanyProfileDto {
        symbol: "API.TEST".to_string(),
        name: "API Test Corp".to_string(),
        industry: Some("API Testing".to_string()),
        list_date: Some(chrono::NaiveDate::from_ymd_opt(2025, 1, 1).unwrap()),
        additional_info: None,
        updated_at: None,
    };

    let request = Request::builder()
        .method("PUT")
        .uri("/api/v1/companies")
        .header("content-type", "application/json")
        .body(Body::from(serde_json::to_string(&new_company).unwrap()))
        .unwrap();

    let response = ServiceExt::oneshot(app.clone().into_service(), request).await.unwrap();
    assert_eq!(response.status(), StatusCode::OK);

    // 2. Act: Get the company
    let request_get = Request::builder()
        .method("GET")
        .uri("/api/v1/companies/API.TEST")
        .body(Body::empty())
        .unwrap();
    
    let response_get = ServiceExt::oneshot(app.clone().into_service(), request_get).await.unwrap();
    assert_eq!(response_get.status(), StatusCode::OK);

    // 3. Assert: Check the response body
    let body = response_get.into_body().collect().await.unwrap().to_bytes();
    let fetched_company: CompanyProfileDto = serde_json::from_slice(&body).unwrap();
    
    assert_eq!(fetched_company.symbol, new_company.symbol);
    assert_eq!(fetched_company.name, new_company.name);
}

#[sqlx::test]
async fn test_api_batch_insert_and_get_financials(pool: PgPool) {
    let state = AppState::new(pool);
    let openapi = app::build_openapi_spec();
    let app = app::build_rest_router_with_state(openapi, state).unwrap();

    // 1. Act: Batch insert financials
    let financials = TimeSeriesFinancialBatchDto {
        records: vec![
            TimeSeriesFinancialDto {
                symbol: "API.FIN".to_string(),
                metric_name: "revenue".to_string(),
                period_date: chrono::NaiveDate::from_ymd_opt(2023, 12, 31).unwrap(),
                value: 2000.0,
                source: Some("api_test".to_string()),
            },
        ],
    };

    let request = Request::builder()
        .method("POST")
        .uri("/api/v1/market-data/financials/batch")
        .header("content-type", "application/json")
        .body(Body::from(serde_json::to_string(&financials).unwrap()))
        .unwrap();

    let response = ServiceExt::oneshot(app.clone().into_service(), request).await.unwrap();
    assert_eq!(response.status(), StatusCode::CREATED);

    // 2. Act: Get the financials
    let request_get = Request::builder()
        .method("GET")
        .uri("/api/v1/market-data/financials/API.FIN?metrics=revenue")
        .body(Body::empty())
        .unwrap();

    let response_get = ServiceExt::oneshot(app.clone().into_service(), request_get).await.unwrap();
    assert_eq!(response_get.status(), StatusCode::OK);

    // 3. Assert: Check the response body
    let body = response_get.into_body().collect().await.unwrap().to_bytes();
    let fetched_financials: Vec<TimeSeriesFinancialDto> = serde_json::from_slice(&body).unwrap();
    
    assert_eq!(fetched_financials.len(), 1);
    assert_eq!(fetched_financials[0].symbol, "API.FIN");
    assert_eq!(fetched_financials[0].metric_name, "revenue");
    assert_eq!(fetched_financials[0].value, 2000.0);
}

#[sqlx::test]
async fn test_api_batch_insert_and_get_daily(pool: PgPool) {
    let state = AppState::new(pool);
    let openapi = app::build_openapi_spec();
    let app = app::build_rest_router_with_state(openapi, state).unwrap();

    // 1. Act: Batch insert daily data
    let daily_data = DailyMarketDataBatchDto {
        records: vec![
            DailyMarketDataDto {
                symbol: "API.DAILY".to_string(),
                trade_date: chrono::NaiveDate::from_ymd_opt(2024, 1, 5).unwrap(),
                close_price: Some(250.5),
                // ... other fields are None
                open_price: None,
                high_price: None,
                low_price: None,
                volume: None,
                pe: None,
                pb: None,
                total_mv: None,
            },
        ],
    };

    let request = Request::builder()
        .method("POST")
        .uri("/api/v1/market-data/daily/batch")
        .header("content-type", "application/json")
        .body(Body::from(serde_json::to_string(&daily_data).unwrap()))
        .unwrap();
    
    let response = ServiceExt::oneshot(app.clone().into_service(), request).await.unwrap();
    assert_eq!(response.status(), StatusCode::CREATED);

    // 2. Act: Get the daily data
    let request_get = Request::builder()
        .method("GET")
        .uri("/api/v1/market-data/daily/API.DAILY")
        .body(Body::empty())
        .unwrap();
    
    let response_get = ServiceExt::oneshot(app.clone().into_service(), request_get).await.unwrap();
    assert_eq!(response_get.status(), StatusCode::OK);

    // 3. Assert: Check the response body
    let body = response_get.into_body().collect().await.unwrap().to_bytes();
    let fetched_data: Vec<DailyMarketDataDto> = serde_json::from_slice(&body).unwrap();
    
    assert_eq!(fetched_data.len(), 1);
    assert_eq!(fetched_data[0].symbol, "API.DAILY");
    assert_eq!(fetched_data[0].close_price, Some(250.5));
}

#[sqlx::test]
async fn test_api_create_and_get_analysis(pool: PgPool) {
    let state = AppState::new(pool);
    let openapi = app::build_openapi_spec();
    let app = app::build_rest_router_with_state(openapi, state).unwrap();

    // 1. Act: Create a new analysis result
    let new_analysis = app::dtos::NewAnalysisResult {
        request_id: Uuid::new_v4(),
        symbol: "API.AI".to_string(),
        template_id: "api_template".to_string(),
        module_id: "bull_case".to_string(),
        content: "This is a test analysis from an API test.".to_string(),
        meta_data: serde_json::json!({ "model": "test-gpt" }),
    };

    let request = Request::builder()
        .method("POST")
        .uri("/api/v1/analysis-results")
        .header("content-type", "application/json")
        .body(Body::from(serde_json::to_string(&new_analysis).unwrap()))
        .unwrap();

    let response = ServiceExt::oneshot(app.clone().into_service(), request).await.unwrap();
    assert_eq!(response.status(), StatusCode::CREATED);

    let body = response.into_body().collect().await.unwrap().to_bytes();
    let created_analysis: app::dtos::AnalysisResultDto = serde_json::from_slice(&body).unwrap();

    // 2. Act: Get the analysis by ID
    let request_get = Request::builder()
        .method("GET")
        .uri(format!("/api/v1/analysis-results/{}", created_analysis.id))
        .body(Body::empty())
        .unwrap();

    let response_get = ServiceExt::oneshot(app.clone().into_service(), request_get).await.unwrap();
    assert_eq!(response_get.status(), StatusCode::OK);
    
    // 3. Assert: Check the response body
    let body_get = response_get.into_body().collect().await.unwrap().to_bytes();
    let fetched_analysis: app::dtos::AnalysisResultDto = serde_json::from_slice(&body_get).unwrap();
    
    assert_eq!(fetched_analysis.id, created_analysis.id);
    assert_eq!(fetched_analysis.symbol, "API.AI");

    // 4. Act: Get by query params
    let request_query = Request::builder()
        .method("GET")
        .uri("/api/v1/analysis-results?symbol=API.AI&module_id=bull_case")
        .body(Body::empty())
        .unwrap();
    
    let response_query = ServiceExt::oneshot(app.clone().into_service(), request_query).await.unwrap();
    assert_eq!(response_query.status(), StatusCode::OK);
    let body_query = response_query.into_body().collect().await.unwrap().to_bytes();
    let fetched_list: Vec<app::dtos::AnalysisResultDto> = serde_json::from_slice(&body_query).unwrap();
    assert_eq!(fetched_list.len(), 1);
    assert_eq!(fetched_list[0].id, created_analysis.id);
}