use serde_json::json;
use rag_module::RagModule;
use std::path::Path;
use tokio::fs;
#[tokio::main]
async fn main() -> anyhow::Result<()> {
println!("🧪 Starting Comprehensive AWS Estate Service Test");
let test_dir = "./comprehensive_test_data";
if Path::new(test_dir).exists() {
fs::remove_dir_all(test_dir).await?;
}
let rag = RagModule::new(test_dir).await?;
rag.initialize().await?;
rag.set_user_context("test-user-123").await?;
println!("✅ RAG Module initialized");
let aws_estate_data = json!([
{
"account_id": "123456789012",
"role_arn": null,
"account_name": "test-account",
"scan_timestamp": "2024-11-05T12:00:00Z",
"services": {
"ec2": {
"instances": [
{
"instance_id": "i-1234567890abcdef0",
"instance_type": "t3.medium",
"state": "running",
"region": "us-west-2",
"launch_time": "2024-11-01T10:00:00Z",
"name": "web-server-1",
"availability_zone": "us-west-2a",
"security_groups": ["sg-12345", "sg-67890"],
"permissions": [
"ec2:StartInstances",
"ec2:StopInstances",
"ec2:DescribeInstances",
"ec2:RebootInstances"
]
},
{
"instance_id": "i-0987654321fedcba0",
"instance_type": "c5.large",
"state": "stopped",
"region": "us-west-2",
"launch_time": "2024-10-15T08:30:00Z",
"name": "batch-processor",
"availability_zone": "us-west-2b",
"permissions": [
"ec2:StartInstances",
"ec2:StopInstances"
]
}
],
"volumes_count": 15,
"total_size_gb": 500,
"attached_volumes": 10,
"unattached_volumes": 5,
"vpcs_count": 3,
"total_subnets": 12,
"internet_gateways_count": 2,
"elastic_ips_total": 5,
"elastic_ips_attached": 3
},
"rds": {
"instances": [
{
"instance_id": "db-instance-1",
"db_instance_identifier": "prod-database",
"region": "us-west-2",
"db_instance_class": "db.r5.large",
"engine": "postgresql",
"db_instance_status": "available",
"name": "production-db",
"multi_az": true,
"allocated_storage": 100,
"storage_type": "gp2",
"permissions": [
"rds:DescribeDBInstances",
"rds:RebootDBInstance"
]
}
],
"clusters_count": 2,
"snapshots_count": 10
},
"s3": {
"buckets": [
{
"bucket_name": "my-app-data-bucket",
"region": "us-west-2",
"creation_date": "2024-01-15T00:00:00Z",
"versioning": "Enabled",
"encryption": "AES256",
"public_access_blocked": true,
"object_count": 1250,
"size_bytes": 52428800,
"permissions": [
"s3:GetObject",
"s3:PutObject",
"s3:ListBucket"
]
},
{
"bucket_name": "backup-storage-bucket",
"region": "us-east-1",
"creation_date": "2024-02-01T00:00:00Z",
"versioning": "Disabled",
"encryption": "Disabled",
"public_access_blocked": true,
"permissions": [
"s3:GetObject",
"s3:ListBucket"
]
}
],
"bucket_policies_count": 3,
"public_bucket_policies": 1
},
"lambda": {
"functions": [
{
"function_name": "data-processor",
"region": "us-west-2",
"runtime": "python3.9",
"timeout": 300,
"memory_size": 1,
"handler": "lambda_function.lambda_handler",
"last_modified": "2024-11-01T12:00:00Z",
"code_size": 1024000,
"role": "arn:aws:iam::123456789012:role/lambda-execution-role",
"environment": {
"Variables": {
"ENV": "production",
"DB_HOST": "prod-database.region.rds.amazonaws.com"
}
},
"permissions": [
"lambda:InvokeFunction",
"lambda:UpdateFunctionCode"
]
}
],
"layers_count": 5,
"event_sources_count": 3
},
"ebs": {
"volumes": [
{
"volume_id": "vol-1234567890abcdef0",
"region": "us-west-2",
"size": 100,
"volume_type": "gp3",
"state": "in-use",
"availability_zone": "us-west-2a",
"encrypted": true,
"create_time": "2024-11-01T10:00:00Z",
"attachment_state": "attached",
"attached_instance_id": "i-1234567890abcdef0",
"iops": 3000,
"throughput": 125,
"permissions": [
"ec2:AttachVolume",
"ec2:DetachVolume"
]
}
],
"snapshots_count": 25,
"snapshots_total_size_gb": 1000,
"encrypted_snapshots": 20,
"total_volumes": 15,
"encrypted_volumes": 12
}
}
}
]);
println!("📊 Test Data Summary:");
println!(" • 1 AWS Account");
println!(" • 5 AWS Services (EC2, RDS, S3, Lambda, EBS)");
let expected_documents = calculate_expected_documents(&aws_estate_data);
println!(" • Expected Documents: {}", expected_documents);
println!("\n📥 Storing AWS estate data...");
let store_result = rag.store_aws_estate_data(aws_estate_data).await?;
println!("✅ Storage completed!");
println!(" • Total accounts processed: {}", store_result.total_accounts);
println!(" • Total services processed: {}", store_result.total_services);
println!(" • Total resources found: {}", store_result.total_resources);
println!(" • Successfully parsed: {}", store_result.parsed_resources);
println!(" • Failed to parse: {}", store_result.failed_resources);
println!(" • Documents created: {}", store_result.create_result.created);
println!(" • Supported services: {:?}", store_result.supported_services);
println!("\n📈 Document Count Analysis:");
println!(" • Expected: {}", expected_documents);
println!(" • Actually created: {}", store_result.create_result.created);
if store_result.create_result.created == expected_documents {
println!(" ✅ Document count MATCHES expectations!");
} else {
println!(" ❌ Document count MISMATCH!");
println!(" • Difference: {}",
(expected_documents as i32) - (store_result.create_result.created as i32));
}
println!("\n📁 Checking local file storage...");
verify_local_storage(test_dir, "test-user-123").await?;
println!("\n🔍 Testing search functionality...");
println!("\n Testing: 'list my ec2 instances'");
let ec2_results = rag.search_estate_resources(
"list my ec2 instances",
None,
None,
None
).await?;
println!(" 📋 EC2 Search Results:");
println!(" • Found {} results", ec2_results.len());
for (i, result) in ec2_results.iter().enumerate() {
if let Some(payload) = &result.payload {
println!(" Result {}: Score: {:.4}", i + 1, result.score);
println!(" • ID: {}", result.id);
println!(" • Instance ID: {}",
payload.get("instance_id")
.and_then(|v| v.as_str())
.unwrap_or("N/A"));
println!(" • IAM Permissions: {}",
payload.get("iam_permissions")
.map(|v| serde_json::to_string_pretty(v).unwrap_or("N/A".to_string()))
.unwrap_or("N/A".to_string()));
}
}
println!("\n Testing: 'show s3 buckets'");
let s3_results = rag.search_estate_resources(
"show s3 buckets",
None,
None,
None
).await?;
println!(" 📋 S3 Search Results:");
println!(" • Found {} results", s3_results.len());
for (i, result) in s3_results.iter().enumerate() {
if let Some(payload) = &result.payload {
println!(" Result {}: Score: {:.4}", i + 1, result.score);
println!(" • Bucket: {}",
payload.get("bucket_name")
.and_then(|v| v.as_str())
.unwrap_or("N/A"));
println!(" • Encryption: {}",
payload.get("encryption")
.and_then(|v| v.as_str())
.unwrap_or("N/A"));
}
}
println!("\n Testing: 'lambda functions'");
let lambda_results = rag.search_estate_resources(
"lambda functions",
None,
None,
None
).await?;
println!(" 📋 Lambda Search Results:");
println!(" • Found {} results", lambda_results.len());
println!("\n🎉 Comprehensive test completed!");
print_final_summary(&store_result, expected_documents, &ec2_results, &s3_results, &lambda_results);
Ok(())
}
fn calculate_expected_documents(data: &serde_json::Value) -> usize {
let mut count = 0;
if let Some(accounts) = data.as_array() {
for account in accounts {
if let Some(services) = account["services"].as_object() {
if let Some(ec2) = services.get("ec2") {
if let Some(instances) = ec2["instances"].as_array() {
count += instances.len();
}
if ec2.get("volumes_count").is_some() {
count += 1;
}
if ec2.get("vpcs_count").is_some() {
count += 1;
}
}
if let Some(rds) = services.get("rds") {
if let Some(instances) = rds["instances"].as_array() {
count += instances.len();
}
if rds["clusters_count"].as_u64().unwrap_or(0) > 0 {
count += 1;
}
if rds["snapshots_count"].as_u64().unwrap_or(0) > 0 {
count += 1;
}
}
if let Some(s3) = services.get("s3") {
if let Some(buckets) = s3["buckets"].as_array() {
count += buckets.len();
}
if s3["bucket_policies_count"].as_u64().unwrap_or(0) > 0 {
count += 1;
}
}
if let Some(lambda) = services.get("lambda") {
if let Some(functions) = lambda["functions"].as_array() {
count += functions.len();
}
if lambda["layers_count"].as_u64().unwrap_or(0) > 0 {
count += 1;
}
if lambda["event_sources_count"].as_u64().unwrap_or(0) > 0 {
count += 1;
}
}
if let Some(ebs) = services.get("ebs") {
if let Some(volumes) = ebs["volumes"].as_array() {
count += volumes.len();
}
if ebs["snapshots_count"].as_u64().unwrap_or(0) > 0 {
count += 1;
}
if ebs.get("total_volumes").is_some() {
count += 1;
}
}
}
}
}
count
}
async fn verify_local_storage(test_dir: &str, user_id: &str) -> anyhow::Result<()> {
let qdrant_data_dir = Path::new(test_dir).join("qdrant-data").join(user_id);
if !qdrant_data_dir.exists() {
println!(" ❌ User directory not found: {:?}", qdrant_data_dir);
return Ok(());
}
println!(" ✅ User directory exists: {:?}", qdrant_data_dir);
let documents_file = qdrant_data_dir.join("aws_estate-documents.json");
let metadata_file = qdrant_data_dir.join("aws_estate-metadata.json");
let vectors_file = qdrant_data_dir.join("aws_estate-vectors.bin");
let index_file = qdrant_data_dir.join("aws_estate-vector-index.json");
println!(" 📄 File existence check:");
println!(" • Documents file: {}", documents_file.exists());
println!(" • Metadata file: {}", metadata_file.exists());
println!(" • Vectors file: {}", vectors_file.exists());
println!(" • Index file: {}", index_file.exists());
if documents_file.exists() {
let content = fs::read_to_string(&documents_file).await?;
if let Ok(wrapper) = serde_json::from_str::<serde_json::Value>(&content) {
if let Some(docs) = wrapper["documents"].as_array() {
println!(" • Documents in local file: {}", docs.len());
}
if let Some(count) = wrapper["count"].as_u64() {
println!(" • Count field in wrapper: {}", count);
}
}
}
Ok(())
}
fn print_final_summary(
store_result: &rag_module::types::AwsEstateIngestResult,
expected_docs: usize,
ec2_results: &[rag_module::types::SearchResult],
s3_results: &[rag_module::types::SearchResult],
lambda_results: &[rag_module::types::SearchResult],
) {
println!("\n" + "=".repeat(60).as_str());
println!("📊 COMPREHENSIVE TEST SUMMARY");
println!("=".repeat(60));
println!("\n🏗️ DOCUMENT CREATION:");
println!(" • Expected documents: {}", expected_docs);
println!(" • Actually created: {}", store_result.create_result.created);
println!(" • Match: {}", if store_result.create_result.created == expected_docs { "✅ YES" } else { "❌ NO" });
println!("\n🔧 SERVICE PARSING:");
println!(" • Supported services: {}", store_result.supported_services.len());
println!(" • Services: {:?}", store_result.supported_services);
println!("\n🔍 SEARCH TESTING:");
println!(" • EC2 search results: {} documents", ec2_results.len());
println!(" • S3 search results: {} documents", s3_results.len());
println!(" • Lambda search results: {} documents", lambda_results.len());
println!("\n🎯 OVERALL STATUS:");
if store_result.create_result.created == expected_docs &&
!ec2_results.is_empty() &&
!s3_results.is_empty() {
println!(" ✅ ALL TESTS PASSED - System working correctly!");
} else {
println!(" ❌ SOME TESTS FAILED - Issues detected!");
}
println!("\n" + "=".repeat(60).as_str());
}