rustytorch/crates/data/rtx-feature-store/tests/integration_tests.rs

//! Integration tests for rtx-feature-store
//!
//! These tests verify the complete functionality of the feature store
//! with real storage backends and comprehensive scenarios.

use chrono::Utc;
use std::collections::HashMap;
use std::time::Duration;
use uuid::Uuid;

use rtx_feature_store::{
    ComputeConfig, EmbeddingConfig, FeatureComputeEngine, FeatureStore, FeatureStoreConfig,
    FeatureStoreError, FeatureValue,
};

/// Test configuration for integration tests
struct TestConfig {
    postgres_url: String,
    redis_url: String,
    vector_store_url: String,
}

impl TestConfig {
    fn new() -> Self {
        Self {
            postgres_url: std::env::var("TEST_POSTGRES_URL").unwrap_or_else(|_| {
                "postgresql://test:test@localhost:5432/test_features".to_string()
            }),
            redis_url: std::env::var("TEST_REDIS_URL")
                .unwrap_or_else(|_| "redis://localhost:6379/1".to_string()),
            vector_store_url: std::env::var("TEST_VECTOR_STORE_URL")
                .unwrap_or_else(|_| "http://localhost:6333".to_string()),
        }
    }
}

#[tokio::test]
#[ignore = "Requires PostgreSQL/Redis infrastructure"]
async fn test_feature_store_basic_operations() {
    let config = TestConfig::new();
    let store_config = FeatureStoreConfig {
        postgres_url: Some(config.postgres_url.clone()),
        redis_url: Some(config.redis_url.clone()),
        vector_store_url: Some(config.vector_store_url.clone()),
        enable_caching: true,
        cache_ttl: Duration::from_secs(300),
        max_connections: 10,
    };

    // Initialize feature store
    let mut store = FeatureStore::new(store_config)
        .await
        .expect("Failed to create feature store");

    // Test feature definition
    let feature_name = "user_embedding";
    let feature_schema = serde_json::json!({
        "type": "array",
        "items": {"type": "number"},
        "minItems": 128,
        "maxItems": 128
    });

    store
        .define_feature(feature_name, feature_schema)
        .await
        .expect("Failed to define feature");

    // Test feature storage
    let user_id = Uuid::new_v4();
    let embedding = vec![0.1f32; 128];
    let feature_value = FeatureValue::Vector(embedding.clone());

    store
        .store_feature(feature_name, &user_id.to_string(), feature_value.clone())
        .await
        .expect("Failed to store feature");

    // Test feature retrieval
    let retrieved = store
        .get_feature(feature_name, &user_id.to_string())
        .await
        .expect("Failed to retrieve feature");

    match retrieved {
        FeatureValue::Vector(vec) => assert_eq!(vec, embedding),
        _ => panic!("Expected vector feature value"),
    }

    // Test batch retrieval
    let entity_ids = vec![user_id.to_string()];
    let batch = store
        .get_features_batch(feature_name, &entity_ids)
        .await
        .expect("Failed to retrieve feature batch");

    assert_eq!(batch.features.len(), 1);
    assert!(batch.features.contains_key(&user_id.to_string()));
}

#[tokio::test]
#[ignore = "Requires PostgreSQL infrastructure"]
async fn test_temporal_consistency() {
    let config = TestConfig::new();
    let store_config = FeatureStoreConfig {
        postgres_url: Some(config.postgres_url.clone()),
        redis_url: None, // Disable caching for temporal consistency test
        vector_store_url: None,
        enable_caching: false,
        cache_ttl: Duration::from_secs(0),
        max_connections: 5,
    };

    let mut store = FeatureStore::new(store_config)
        .await
        .expect("Failed to create feature store");

    let feature_name = "user_score";
    let entity_id = "user_123";

    // Define feature
    let feature_schema = serde_json::json!({"type": "number"});
    store
        .define_feature(feature_name, feature_schema)
        .await
        .expect("Failed to define feature");

    // Store features at different timestamps
    let t1 = Utc::now();
    store
        .store_feature_at_time(feature_name, entity_id, FeatureValue::Float(0.5), t1)
        .await
        .expect("Failed to store feature at t1");

    tokio::time::sleep(Duration::from_millis(100)).await;

    let t2 = Utc::now();
    store
        .store_feature_at_time(feature_name, entity_id, FeatureValue::Float(0.8), t2)
        .await
        .expect("Failed to store feature at t2");

    tokio::time::sleep(Duration::from_millis(100)).await;

    let t3 = Utc::now();
    store
        .store_feature_at_time(feature_name, entity_id, FeatureValue::Float(1.0), t3)
        .await
        .expect("Failed to store feature at t3");

    // Test point-in-time retrieval
    let value_at_t1 = store
        .get_feature_at_time(feature_name, entity_id, t1)
        .await
        .expect("Failed to retrieve feature at t1");
    assert_eq!(value_at_t1, FeatureValue::Float(0.5));

    let value_at_t2 = store
        .get_feature_at_time(feature_name, entity_id, t2)
        .await
        .expect("Failed to retrieve feature at t2");
    assert_eq!(value_at_t2, FeatureValue::Float(0.8));

    let value_at_t3 = store
        .get_feature_at_time(feature_name, entity_id, t3)
        .await
        .expect("Failed to retrieve feature at t3");
    assert_eq!(value_at_t3, FeatureValue::Float(1.0));

    // Test historical query (should return value at t2 when querying between t2 and t3)
    let t_between = t2 + chrono::Duration::milliseconds(50);
    let value_between = store
        .get_feature_at_time(feature_name, entity_id, t_between)
        .await
        .expect("Failed to retrieve feature at t_between");
    assert_eq!(value_between, FeatureValue::Float(0.8));
}

#[tokio::test]
#[ignore = "Requires PostgreSQL/Redis infrastructure"]
async fn test_online_serving_latency() {
    let config = TestConfig::new();
    let store_config = FeatureStoreConfig {
        postgres_url: Some(config.postgres_url.clone()),
        redis_url: Some(config.redis_url.clone()),
        vector_store_url: None,
        enable_caching: true,
        cache_ttl: Duration::from_secs(300),
        max_connections: 10,
    };

    let mut store = FeatureStore::new(store_config)
        .await
        .expect("Failed to create feature store");

    // Setup serving configuration
    // Test basic feature storage operations instead of full serving
    // (Full serving tests would require running server infrastructure)

    // Define test features
    store
        .define_feature("feature_1", serde_json::json!({"type": "number"}))
        .await
        .expect("Failed to define feature_1");
    store
        .define_feature("feature_2", serde_json::json!({"type": "number"}))
        .await
        .expect("Failed to define feature_2");
    store
        .define_feature("feature_3", serde_json::json!({"type": "number"}))
        .await
        .expect("Failed to define feature_3");

    let entity_id = "test_entity";

    // Store features
    store
        .store_feature("feature_1", entity_id, FeatureValue::Float(0.0))
        .await
        .expect("Failed to store feature_1");
    store
        .store_feature("feature_2", entity_id, FeatureValue::Float(1.0))
        .await
        .expect("Failed to store feature_2");
    store
        .store_feature("feature_3", entity_id, FeatureValue::Float(2.0))
        .await
        .expect("Failed to store feature_3");

    // Verify retrieval latency
    let start = std::time::Instant::now();
    let value = store
        .get_feature("feature_1", entity_id)
        .await
        .expect("Failed to retrieve feature");
    let latency = start.elapsed();

    assert!(
        latency < Duration::from_millis(100),
        "Latency too high: {:?}",
        latency
    );
    assert_eq!(value, FeatureValue::Float(0.0));
}

#[tokio::test]
#[ignore = "Requires PostgreSQL infrastructure"]
async fn test_feature_versioning() {
    let config = TestConfig::new();
    let store_config = FeatureStoreConfig {
        postgres_url: Some(config.postgres_url.clone()),
        redis_url: None,
        vector_store_url: None,
        enable_caching: false,
        cache_ttl: Duration::from_secs(0),
        max_connections: 5,
    };

    let mut store = FeatureStore::new(store_config)
        .await
        .expect("Failed to create feature store");

    let feature_name = "versioned_feature";
    let entity_id = "entity_123";

    // Create different versions of the same feature
    let v1_schema = serde_json::json!({"type": "number"});
    let v1_version = store
        .create_feature_version(feature_name, v1_schema, "Initial version")
        .await
        .expect("Failed to create v1");

    let v2_schema = serde_json::json!({
        "type": "object",
        "properties": {"score": {"type": "number"}, "confidence": {"type": "number"}}
    });
    let v2_version = store
        .create_feature_version(feature_name, v2_schema, "Added confidence")
        .await
        .expect("Failed to create v2");

    // Store data for each version
    store
        .store_feature_version(
            feature_name,
            entity_id,
            FeatureValue::Float(0.5),
            v1_version,
        )
        .await
        .expect("Failed to store v1 feature");

    let v2_data = serde_json::json!({"score": 0.8, "confidence": 0.9});
    store
        .store_feature_version(
            feature_name,
            entity_id,
            FeatureValue::Object(v2_data.clone()),
            v2_version,
        )
        .await
        .expect("Failed to store v2 feature");

    // Test version-specific retrieval
    let v1_retrieved = store
        .get_feature_version(feature_name, entity_id, v1_version)
        .await
        .expect("Failed to retrieve v1 feature");
    assert_eq!(v1_retrieved, FeatureValue::Float(0.5));

    let v2_retrieved = store
        .get_feature_version(feature_name, entity_id, v2_version)
        .await
        .expect("Failed to retrieve v2 feature");
    assert_eq!(v2_retrieved, FeatureValue::Object(v2_data));

    // Test lineage tracking
    let lineage = store
        .get_feature_lineage(feature_name)
        .await
        .expect("Failed to get feature lineage");
    assert_eq!(lineage.versions.len(), 2);
    assert_eq!(lineage.versions[0].version_id, v1_version);
    assert_eq!(lineage.versions[1].version_id, v2_version);
}

#[tokio::test]
#[ignore = "Requires PostgreSQL/Redis infrastructure"]
async fn test_feature_drift_monitoring() {
    let config = TestConfig::new();
    let store_config = FeatureStoreConfig {
        postgres_url: Some(config.postgres_url.clone()),
        redis_url: Some(config.redis_url.clone()),
        vector_store_url: None,
        enable_caching: true,
        cache_ttl: Duration::from_secs(300),
        max_connections: 10,
    };

    let mut store = FeatureStore::new(store_config)
        .await
        .expect("Failed to create feature store");

    let feature_name = "drift_monitored_feature";
    let feature_schema = serde_json::json!({"type": "number"});

    store
        .define_feature(feature_name, feature_schema)
        .await
        .expect("Failed to define feature");

    // Enable drift monitoring for this feature
    store
        .enable_drift_monitoring(feature_name, 0.1)
        .await // 10% drift threshold
        .expect("Failed to enable drift monitoring");

    // Store baseline data
    let baseline_entities = (0..100)
        .map(|i| format!("entity_{}", i))
        .collect::<Vec<_>>();
    for (i, entity_id) in baseline_entities.iter().enumerate() {
        let value = FeatureValue::Float(i as f64 / 100.0); // Values 0.0 to 0.99
        store
            .store_feature(feature_name, entity_id, value)
            .await
            .expect("Failed to store baseline feature");
    }

    // Wait for baseline calculation
    tokio::time::sleep(Duration::from_millis(100)).await;

    // Store drifted data (shifted distribution)
    let drift_entities = (100..200)
        .map(|i| format!("entity_{}", i))
        .collect::<Vec<_>>();
    for (i, entity_id) in drift_entities.iter().enumerate() {
        let value = FeatureValue::Float((i as f64 + 50.0) / 100.0); // Values 0.5 to 1.49 (shifted)
        store
            .store_feature(feature_name, entity_id, value)
            .await
            .expect("Failed to store drifted feature");
    }

    // Check for drift detection
    tokio::time::sleep(Duration::from_millis(500)).await;

    let drift_report = store
        .get_drift_report(feature_name)
        .await
        .expect("Failed to get drift report");

    assert!(
        drift_report.drift_detected,
        "Drift should have been detected"
    );
    assert!(
        drift_report.drift_score > 0.1,
        "Drift score should exceed threshold"
    );
    assert!(
        !drift_report.alerts.is_empty(),
        "Should have generated alerts"
    );
}

#[tokio::test]
#[ignore = "Requires PostgreSQL/Redis/Vector Store infrastructure"]
async fn test_vector_store_integration() {
    let config = TestConfig::new();
    let store_config = FeatureStoreConfig {
        postgres_url: Some(config.postgres_url.clone()),
        redis_url: Some(config.redis_url.clone()),
        vector_store_url: Some(config.vector_store_url.clone()),
        enable_caching: true,
        cache_ttl: Duration::from_secs(300),
        max_connections: 10,
    };

    let mut store = FeatureStore::new(store_config)
        .await
        .expect("Failed to create feature store");

    let collection_name = "test_embeddings";
    let embedding_config = EmbeddingConfig {
        dimension: 128,
        distance_metric: "cosine".to_string(),
        index_type: "hnsw".to_string(),
    };

    // Create vector collection
    store
        .create_vector_collection(collection_name, embedding_config)
        .await
        .expect("Failed to create vector collection");

    // Store embeddings
    let entities = (0..10).map(|i| format!("item_{}", i)).collect::<Vec<_>>();
    let embeddings = entities
        .iter()
        .enumerate()
        .map(|(i, entity_id)| {
            let mut embedding = vec![0.0f32; 128];
            embedding[i % 128] = 1.0; // Unique embedding pattern
            (entity_id.clone(), embedding)
        })
        .collect::<Vec<_>>();

    for (entity_id, embedding) in &embeddings {
        store
            .store_vector(collection_name, entity_id, embedding.clone())
            .await
            .expect("Failed to store vector");
    }

    // Test similarity search
    let query_embedding = embeddings[0].1.clone();
    let similar_items = store
        .search_similar_vectors(collection_name, &query_embedding, 5)
        .await
        .expect("Failed to search similar vectors");

    assert!(!similar_items.is_empty(), "Should find similar vectors");
    assert_eq!(
        similar_items[0].entity_id, entities[0],
        "Most similar should be the same item"
    );
    assert!(
        similar_items[0].score > 0.9,
        "Similarity score should be high for identical vector"
    );

    // Test batch vector retrieval
    let batch_entities = vec![entities[0].clone(), entities[1].clone()];
    let batch_vectors = store
        .get_vectors_batch(collection_name, &batch_entities)
        .await
        .expect("Failed to retrieve vector batch");

    assert_eq!(batch_vectors.len(), 2);
    assert_eq!(batch_vectors[&entities[0]], embeddings[0].1);
    assert_eq!(batch_vectors[&entities[1]], embeddings[1].1);
}

#[tokio::test]
#[ignore = "Requires PostgreSQL/Redis infrastructure"]
async fn test_compute_engine() {
    let config = TestConfig::new();
    let store_config = FeatureStoreConfig {
        postgres_url: Some(config.postgres_url.clone()),
        redis_url: Some(config.redis_url.clone()),
        vector_store_url: None,
        enable_caching: true,
        cache_ttl: Duration::from_secs(300),
        max_connections: 10,
    };

    let mut store = FeatureStore::new(store_config)
        .await
        .expect("Failed to create feature store");

    let compute_config = ComputeConfig {
        max_concurrent_jobs: 5,
        batch_size: 100,
        retry_attempts: 3,
        timeout: Duration::from_secs(30),
    };

    let compute_engine = FeatureComputeEngine::new(compute_config);

    // Define raw features
    store
        .define_feature("raw_clicks", serde_json::json!({"type": "integer"}))
        .await
        .expect("Failed to define raw_clicks");
    store
        .define_feature("raw_views", serde_json::json!({"type": "integer"}))
        .await
        .expect("Failed to define raw_views");

    // Store raw data
    let entity_ids = (0..10).map(|i| format!("user_{}", i)).collect::<Vec<_>>();
    for (i, entity_id) in entity_ids.iter().enumerate() {
        store
            .store_feature(
                "raw_clicks",
                entity_id,
                FeatureValue::Integer((i as i64 + 1) * 10),
            )
            .await
            .expect("Failed to store raw_clicks");
        store
            .store_feature(
                "raw_views",
                entity_id,
                FeatureValue::Integer((i as i64 + 1) * 100),
            )
            .await
            .expect("Failed to store raw_views");
    }

    // Define computed feature
    let ctr_computation = |inputs: &HashMap<String, FeatureValue>| -> Result<FeatureValue, Box<dyn std::error::Error>> {
        let clicks = match inputs.get("raw_clicks") {
            Some(FeatureValue::Integer(c)) => *c as f64,
            _ => return Err("Missing or invalid clicks".into()),
        };
        let views = match inputs.get("raw_views") {
            Some(FeatureValue::Integer(v)) => *v as f64,
            _ => return Err("Missing or invalid views".into()),
        };

        let ctr = if views > 0.0 { clicks / views } else { 0.0 };
        Ok(FeatureValue::Float(ctr))
    };

    store
        .define_computed_feature(
            "click_through_rate",
            vec!["raw_clicks".to_string(), "raw_views".to_string()],
            Box::new(ctr_computation),
        )
        .await
        .expect("Failed to define computed feature");

    // Trigger computation
    compute_engine
        .compute_features(&mut store, &entity_ids)
        .await
        .expect("Failed to compute features");

    // Verify computed features
    for entity_id in &entity_ids {
        let ctr = store
            .get_feature("click_through_rate", entity_id)
            .await
            .expect("Failed to retrieve computed feature");

        match ctr {
            FeatureValue::Float(value) => assert_eq!(value, 0.1), // clicks/views = 10/100 = 0.1
            _ => panic!("Expected float CTR value"),
        }
    }
}

#[tokio::test]
#[ignore = "Requires PostgreSQL infrastructure"]
async fn test_error_handling() {
    let config = TestConfig::new();
    let store_config = FeatureStoreConfig {
        postgres_url: Some("invalid_url".to_string()), // Invalid URL for testing
        redis_url: None,
        vector_store_url: None,
        enable_caching: false,
        cache_ttl: Duration::from_secs(0),
        max_connections: 1,
    };

    // Should fail to connect with invalid URL
    let store_result = FeatureStore::new(store_config).await;
    assert!(
        store_result.is_err(),
        "Should fail with invalid PostgreSQL URL"
    );

    if let Err(error) = store_result {
        match error {
            FeatureStoreError::Storage(_) => {} // Expected
            _ => panic!("Expected storage error, got: {:?}", error),
        }
    }

    // Test with valid store but invalid operations
    let valid_config = FeatureStoreConfig {
        postgres_url: Some(TestConfig::new().postgres_url),
        redis_url: None,
        vector_store_url: None,
        enable_caching: false,
        cache_ttl: Duration::from_secs(0),
        max_connections: 5,
    };

    let mut store = FeatureStore::new(valid_config)
        .await
        .expect("Failed to create valid store");

    // Test retrieving non-existent feature
    let result = store.get_feature("non_existent_feature", "entity_id").await;
    assert!(result.is_err(), "Should fail for non-existent feature");

    let error = result.unwrap_err();
    match error {
        FeatureStoreError::NotFound(_) => {} // Expected
        _ => panic!("Expected not found error, got: {:?}", error),
    }
}

#[tokio::test]
#[ignore = "Requires PostgreSQL/Redis infrastructure"]
async fn test_concurrent_access() {
    let config = TestConfig::new();
    let store_config = FeatureStoreConfig {
        postgres_url: Some(config.postgres_url.clone()),
        redis_url: Some(config.redis_url.clone()),
        vector_store_url: None,
        enable_caching: true,
        cache_ttl: Duration::from_secs(300),
        max_connections: 20, // Higher connection limit for concurrency
    };

    let store = std::sync::Arc::new(tokio::sync::Mutex::new(
        FeatureStore::new(store_config)
            .await
            .expect("Failed to create feature store"),
    ));

    // Define test feature
    {
        let mut store_lock = store.lock().await;
        store_lock
            .define_feature("concurrent_feature", serde_json::json!({"type": "number"}))
            .await
            .expect("Failed to define feature");
    }

    // Launch concurrent read/write operations
    let mut handles = vec![];
    for i in 0..50 {
        let store_clone = store.clone();
        let entity_id = format!("entity_{}", i);
        let handle = tokio::spawn(async move {
            let mut store_lock = store_clone.lock().await;

            // Write operation
            let value = FeatureValue::Float(i as f64);
            store_lock
                .store_feature("concurrent_feature", &entity_id, value.clone())
                .await
                .expect("Failed to store feature concurrently");

            // Read operation
            let retrieved = store_lock
                .get_feature("concurrent_feature", &entity_id)
                .await
                .expect("Failed to retrieve feature concurrently");

            assert_eq!(
                retrieved, value,
                "Retrieved value should match stored value"
            );
        });
        handles.push(handle);
    }

    // Wait for all operations to complete
    for handle in handles {
        handle.await.expect("Concurrent operation failed");
    }

    // Verify all data was stored correctly
    let mut store_lock = store.lock().await;
    for i in 0..50 {
        let entity_id = format!("entity_{}", i);
        let value = store_lock
            .get_feature("concurrent_feature", &entity_id)
            .await
            .expect("Failed to verify concurrent data");
        assert_eq!(value, FeatureValue::Float(i as f64));
    }
}