use aimo_core::receipt::RequestReceipt;
use solana_sdk::{signature::Keypair, signer::Signer};
#[test]
fn test_request_receipt_builder_all_fields() {
let sender_keypair = Keypair::new();
let request_id_keypair = Keypair::new();
let receipt = RequestReceipt::new()
.sender(sender_keypair.pubkey())
.request_id(request_id_keypair.pubkey())
.service_id("test-service-123")
.max_est_cost(1000)
.timestamp(1684512000)
.finish_reason("stop")
.prompt_tokens(150)
.completion_tokens(75)
.amount(500)
.token_mint("SOL")
.model_name("gpt-4")
.cache_discount(50)
.streamed(true)
.cancelled(false)
.latency(250)
.generation_time(1200)
.moderation_latency(50)
.num_media_prompt(2)
.num_media_completion(1)
.num_search_results(5)
.app_id("my-app")
.origin("https://example.com");
assert_eq!(receipt.sender, sender_keypair.pubkey().to_string());
assert_eq!(receipt.request_id, request_id_keypair.pubkey().to_string());
assert_eq!(receipt.service_id, "test-service-123");
assert_eq!(receipt.max_est_cost, 1000);
assert_eq!(receipt.timestamp, 1684512000);
assert_eq!(receipt.finish_reason, "stop");
assert_eq!(receipt.prompt_tokens, 150);
assert_eq!(receipt.completion_tokens, 75);
assert_eq!(receipt.amount, 500);
assert_eq!(receipt.token_mint, "SOL");
assert_eq!(receipt.model_name, "gpt-4");
assert_eq!(receipt.cache_discount, 50);
assert_eq!(receipt.streamed, Some(true));
assert_eq!(receipt.cancelled, Some(false));
assert_eq!(receipt.latency, Some(250));
assert_eq!(receipt.generation_time, Some(1200));
assert_eq!(receipt.moderation_latency, Some(50));
assert_eq!(receipt.num_media_prompt, Some(2));
assert_eq!(receipt.num_media_completion, Some(1));
assert_eq!(receipt.num_search_results, Some(5));
assert_eq!(receipt.app_id, Some("my-app".to_string()));
assert_eq!(receipt.origin, Some("https://example.com".to_string()));
}
#[test]
fn test_request_receipt_builder_partial() {
let sender_keypair = Keypair::new();
let request_id_keypair = Keypair::new();
let receipt = RequestReceipt::new()
.sender(sender_keypair.pubkey())
.request_id(request_id_keypair.pubkey())
.service_id("minimal-service")
.amount(100)
.token_mint("USDC")
.model_name("gpt-3.5-turbo");
assert_eq!(receipt.sender, sender_keypair.pubkey().to_string());
assert_eq!(receipt.request_id, request_id_keypair.pubkey().to_string());
assert_eq!(receipt.service_id, "minimal-service");
assert_eq!(receipt.amount, 100);
assert_eq!(receipt.token_mint, "USDC");
assert_eq!(receipt.model_name, "gpt-3.5-turbo");
assert_eq!(receipt.max_est_cost, 0);
assert_eq!(receipt.timestamp, 0);
assert_eq!(receipt.finish_reason, "");
assert_eq!(receipt.prompt_tokens, 0);
assert_eq!(receipt.completion_tokens, 0);
assert_eq!(receipt.cache_discount, 0);
assert_eq!(receipt.streamed, None);
assert_eq!(receipt.cancelled, None);
assert_eq!(receipt.latency, None);
assert_eq!(receipt.generation_time, None);
assert_eq!(receipt.moderation_latency, None);
assert_eq!(receipt.num_media_prompt, None);
assert_eq!(receipt.num_media_completion, None);
assert_eq!(receipt.num_search_results, None);
assert_eq!(receipt.app_id, None);
assert_eq!(receipt.origin, None);
}
#[test]
fn test_request_receipt_builder_chaining() {
let sender_keypair = Keypair::new();
let request_id_keypair = Keypair::new();
let receipt1 = RequestReceipt::new()
.amount(100)
.sender(sender_keypair.pubkey())
.model_name("test-model")
.request_id(request_id_keypair.pubkey())
.service_id("test-service");
let receipt2 = RequestReceipt::new()
.sender(sender_keypair.pubkey())
.request_id(request_id_keypair.pubkey())
.service_id("test-service")
.model_name("test-model")
.amount(100);
assert_eq!(receipt1.sender, receipt2.sender);
assert_eq!(receipt1.request_id, receipt2.request_id);
assert_eq!(receipt1.service_id, receipt2.service_id);
assert_eq!(receipt1.model_name, receipt2.model_name);
assert_eq!(receipt1.amount, receipt2.amount);
}
#[test]
fn test_request_receipt_builder_into_string_types() {
let sender_keypair = Keypair::new();
let receipt = RequestReceipt::new()
.sender(sender_keypair.pubkey())
.service_id("string_literal")
.finish_reason(String::from("owned_string"))
.token_mint("SOL")
.model_name(format!("formatted_{}", "string"))
.app_id("app_literal")
.origin(String::from("https://test.com"));
assert_eq!(receipt.service_id, "string_literal");
assert_eq!(receipt.finish_reason, "owned_string");
assert_eq!(receipt.token_mint, "SOL");
assert_eq!(receipt.model_name, "formatted_string");
assert_eq!(receipt.app_id, Some("app_literal".to_string()));
assert_eq!(receipt.origin, Some("https://test.com".to_string()));
}
#[test]
fn test_request_receipt_serialization() {
let sender_keypair = Keypair::new();
let request_id_keypair = Keypair::new();
let receipt = RequestReceipt::new()
.sender(sender_keypair.pubkey())
.request_id(request_id_keypair.pubkey())
.service_id("test-service")
.amount(100)
.prompt_tokens(50)
.completion_tokens(25)
.streamed(true)
.app_id("test-app");
let bytes_result = receipt.clone().to_bytes();
assert!(bytes_result.is_ok());
let hash_result = receipt.to_hash();
assert!(hash_result.is_ok());
let hash = hash_result.unwrap();
assert_eq!(hash.len(), 32); }