//! 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::>(); 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::>(); 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::>(); 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::>(); 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::>(); 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| -> Result> { 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)); } }