448 lines
14 KiB
Rust
448 lines
14 KiB
Rust
//! Comprehensive tests for Feature Store functionality
|
|
|
|
use rtx_feature_store::{FeatureStore, FeatureStoreConfig, FeatureValue};
|
|
use std::time::Duration;
|
|
|
|
#[cfg(test)]
|
|
mod feature_store_tests {
|
|
use super::*;
|
|
|
|
fn create_test_config() -> FeatureStoreConfig {
|
|
FeatureStoreConfig {
|
|
postgres_url: None,
|
|
redis_url: None,
|
|
vector_store_url: None,
|
|
enable_caching: false,
|
|
cache_ttl: Duration::from_secs(300),
|
|
max_connections: 10,
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_store_creation() {
|
|
let config = create_test_config();
|
|
let store = FeatureStore::new(config).await;
|
|
assert!(store.is_ok());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_store_and_retrieve_feature() {
|
|
let config = create_test_config();
|
|
let mut store = FeatureStore::new(config).await.unwrap();
|
|
|
|
// Define feature first
|
|
let schema = serde_json::json!({"type": "number"});
|
|
store.define_feature("test_feature", schema).await.unwrap();
|
|
|
|
let feature = FeatureValue::Float(42.5);
|
|
|
|
// Store feature
|
|
let result = store
|
|
.store_feature("test_feature", "entity1", feature.clone())
|
|
.await;
|
|
assert!(result.is_ok());
|
|
|
|
// Retrieve feature
|
|
let retrieved = store.get_feature("test_feature", "entity1").await;
|
|
assert!(retrieved.is_ok());
|
|
assert_eq!(retrieved.unwrap(), feature);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_batch_store_features() {
|
|
let config = create_test_config();
|
|
let mut store = FeatureStore::new(config).await.unwrap();
|
|
|
|
// Define features first
|
|
store
|
|
.define_feature("feature1", serde_json::json!({"type": "number"}))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.define_feature("feature2", serde_json::json!({"type": "integer"}))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.define_feature("feature3", serde_json::json!({"type": "string"}))
|
|
.await
|
|
.unwrap();
|
|
|
|
// Store features individually
|
|
store
|
|
.store_feature("feature1", "entity1", FeatureValue::Float(1.0))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.store_feature("feature2", "entity1", FeatureValue::Integer(42))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.store_feature(
|
|
"feature3",
|
|
"entity1",
|
|
FeatureValue::String("test".to_string()),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Retrieve each feature
|
|
let retrieved1 = store.get_feature("feature1", "entity1").await;
|
|
assert!(retrieved1.is_ok());
|
|
assert_eq!(retrieved1.unwrap(), FeatureValue::Float(1.0));
|
|
|
|
let retrieved2 = store.get_feature("feature2", "entity1").await;
|
|
assert!(retrieved2.is_ok());
|
|
assert_eq!(retrieved2.unwrap(), FeatureValue::Integer(42));
|
|
|
|
let retrieved3 = store.get_feature("feature3", "entity1").await;
|
|
assert!(retrieved3.is_ok());
|
|
assert_eq!(
|
|
retrieved3.unwrap(),
|
|
FeatureValue::String("test".to_string())
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_feature_versioning() {
|
|
let config = FeatureStoreConfig {
|
|
postgres_url: None,
|
|
redis_url: None,
|
|
vector_store_url: None,
|
|
enable_caching: false,
|
|
cache_ttl: Duration::from_secs(300),
|
|
max_connections: 10,
|
|
};
|
|
let mut store = FeatureStore::new(config).await.unwrap();
|
|
|
|
// Define feature
|
|
let schema = serde_json::json!({"type": "number"});
|
|
store
|
|
.define_feature("versioned_feature", schema)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Create version 1
|
|
let v1_version = store
|
|
.create_feature_version(
|
|
"versioned_feature",
|
|
serde_json::json!({"type": "number"}),
|
|
"Version 1",
|
|
)
|
|
.await
|
|
.unwrap();
|
|
let v1 = FeatureValue::Float(1.0);
|
|
store
|
|
.store_feature_version("versioned_feature", "entity1", v1.clone(), v1_version)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Create version 2
|
|
let v2_version = store
|
|
.create_feature_version(
|
|
"versioned_feature",
|
|
serde_json::json!({"type": "number"}),
|
|
"Version 2",
|
|
)
|
|
.await
|
|
.unwrap();
|
|
let v2 = FeatureValue::Float(2.0);
|
|
store
|
|
.store_feature_version("versioned_feature", "entity1", v2.clone(), v2_version)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Latest version should be v2 (last stored)
|
|
let latest = store
|
|
.get_feature("versioned_feature", "entity1")
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(latest, v2);
|
|
|
|
// Should be able to get specific version
|
|
let specific = store
|
|
.get_feature_version("versioned_feature", "entity1", v1_version)
|
|
.await;
|
|
assert!(specific.is_ok());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_feature_expiration() {
|
|
let config = FeatureStoreConfig {
|
|
postgres_url: None,
|
|
redis_url: Some("redis://localhost:6379/9".to_string()), // Test DB
|
|
vector_store_url: None,
|
|
enable_caching: true,
|
|
cache_ttl: Duration::from_secs(1), // 1 second TTL
|
|
max_connections: 10,
|
|
};
|
|
|
|
// Skip test if Redis is not available
|
|
if let Ok(mut store) = FeatureStore::new(config).await {
|
|
let schema = serde_json::json!({"type": "number"});
|
|
store
|
|
.define_feature("expiring_feature", schema)
|
|
.await
|
|
.unwrap();
|
|
|
|
let feature = FeatureValue::Float(99.9);
|
|
store
|
|
.store_feature("expiring_feature", "entity1", feature.clone())
|
|
.await
|
|
.unwrap();
|
|
|
|
// Should retrieve immediately
|
|
let result = store.get_feature("expiring_feature", "entity1").await;
|
|
assert!(result.is_ok());
|
|
assert_eq!(result.unwrap(), feature);
|
|
|
|
// Wait for expiration
|
|
tokio::time::sleep(Duration::from_secs(2)).await;
|
|
|
|
// Cache should be expired (but data still in primary storage)
|
|
// This test mainly verifies no panics occur
|
|
let _ = store.get_feature("expiring_feature", "entity1").await;
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_delete_feature() {
|
|
let config = create_test_config();
|
|
let mut store = FeatureStore::new(config).await.unwrap();
|
|
|
|
let schema = serde_json::json!({"type": "integer"});
|
|
store.define_feature("deletable", schema).await.unwrap();
|
|
|
|
let feature = FeatureValue::Integer(123);
|
|
store
|
|
.store_feature("deletable", "entity1", feature)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Should exist
|
|
assert!(store.get_feature("deletable", "entity1").await.is_ok());
|
|
|
|
// Delete it (deletes the entire feature, not just for one entity)
|
|
let deleted = store.delete_feature("deletable").await;
|
|
assert!(deleted.is_ok());
|
|
|
|
// Should not exist anymore
|
|
let result = store.get_feature("deletable", "entity1").await;
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_list_features() {
|
|
let config = create_test_config();
|
|
let mut store = FeatureStore::new(config).await.unwrap();
|
|
|
|
// Define and add multiple features
|
|
store
|
|
.define_feature("feat1", serde_json::json!({"type": "number"}))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.define_feature("feat2", serde_json::json!({"type": "number"}))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.define_feature("feat3", serde_json::json!({"type": "number"}))
|
|
.await
|
|
.unwrap();
|
|
|
|
store
|
|
.store_feature("feat1", "entity1", FeatureValue::Float(1.0))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.store_feature("feat2", "entity1", FeatureValue::Float(2.0))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.store_feature("feat3", "entity1", FeatureValue::Float(3.0))
|
|
.await
|
|
.unwrap();
|
|
|
|
let features = store.list_features().await;
|
|
assert!(features.is_ok());
|
|
let feature_names = features.unwrap();
|
|
assert_eq!(feature_names.len(), 3);
|
|
assert!(feature_names.contains(&"feat1".to_string()));
|
|
assert!(feature_names.contains(&"feat2".to_string()));
|
|
assert!(feature_names.contains(&"feat3".to_string()));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_feature_value_types() {
|
|
let config = create_test_config();
|
|
let mut store = FeatureStore::new(config).await.unwrap();
|
|
|
|
// Test all value types
|
|
let test_cases = vec![
|
|
("float_feat", FeatureValue::Float(3.14)),
|
|
("int_feat", FeatureValue::Integer(42)),
|
|
("string_feat", FeatureValue::String("hello".to_string())),
|
|
("bool_feat", FeatureValue::Boolean(true)),
|
|
(
|
|
"vector_feat",
|
|
FeatureValue::Vector(vec![1.0, 2.0, 3.0, 4.0]),
|
|
),
|
|
(
|
|
"array_feat",
|
|
FeatureValue::Array(vec![FeatureValue::Float(1.0), FeatureValue::Float(2.0)]),
|
|
),
|
|
];
|
|
|
|
for (name, value) in &test_cases {
|
|
store
|
|
.define_feature(name, serde_json::json!({"type": "any"}))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.store_feature(name, "entity1", value.clone())
|
|
.await
|
|
.unwrap();
|
|
let retrieved = store.get_feature(name, "entity1").await.unwrap();
|
|
assert_eq!(retrieved, *value);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_concurrent_access() {
|
|
let config = create_test_config();
|
|
let store = FeatureStore::new(config).await.unwrap();
|
|
let store = std::sync::Arc::new(tokio::sync::Mutex::new(store));
|
|
|
|
// Define features first
|
|
{
|
|
let mut store_lock = store.lock().await;
|
|
for i in 0..10 {
|
|
let feature_name = format!("concurrent_feat_{}", i);
|
|
store_lock
|
|
.define_feature(&feature_name, serde_json::json!({"type": "number"}))
|
|
.await
|
|
.unwrap();
|
|
}
|
|
}
|
|
|
|
let mut handles = vec![];
|
|
|
|
// Spawn multiple tasks that write and read concurrently
|
|
for i in 0..10 {
|
|
let store_clone = store.clone();
|
|
let handle = tokio::spawn(async move {
|
|
let mut store = store_clone.lock().await;
|
|
let feature_name = format!("concurrent_feat_{}", i);
|
|
let value = FeatureValue::Float(i as f64);
|
|
|
|
// Write
|
|
store
|
|
.store_feature(&feature_name, "entity1", value.clone())
|
|
.await
|
|
.unwrap();
|
|
|
|
// Read back
|
|
let retrieved = store.get_feature(&feature_name, "entity1").await.unwrap();
|
|
assert_eq!(retrieved, value);
|
|
});
|
|
handles.push(handle);
|
|
}
|
|
|
|
// Wait for all tasks to complete
|
|
for handle in handles {
|
|
handle.await.unwrap();
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_feature_search() {
|
|
let config = create_test_config();
|
|
let mut store = FeatureStore::new(config).await.unwrap();
|
|
|
|
// Define and add features with patterns
|
|
store
|
|
.define_feature("user_age", serde_json::json!({"type": "integer"}))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.define_feature("user_name", serde_json::json!({"type": "string"}))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.define_feature("user_score", serde_json::json!({"type": "number"}))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.define_feature("item_price", serde_json::json!({"type": "number"}))
|
|
.await
|
|
.unwrap();
|
|
|
|
store
|
|
.store_feature("user_age", "entity1", FeatureValue::Integer(25))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.store_feature(
|
|
"user_name",
|
|
"entity1",
|
|
FeatureValue::String("Alice".to_string()),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.store_feature("user_score", "entity1", FeatureValue::Float(95.5))
|
|
.await
|
|
.unwrap();
|
|
store
|
|
.store_feature("item_price", "entity1", FeatureValue::Float(19.99))
|
|
.await
|
|
.unwrap();
|
|
|
|
// List all features (search functionality would be a separate feature)
|
|
let all_features = store.list_features().await;
|
|
assert!(all_features.is_ok());
|
|
let results = all_features.unwrap();
|
|
assert_eq!(results.len(), 4);
|
|
|
|
// Verify user features are present
|
|
let user_features: Vec<_> = results.iter().filter(|f| f.starts_with("user_")).collect();
|
|
assert_eq!(user_features.len(), 3);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_redis_backend_operations() {
|
|
// Skip if Redis is not available
|
|
if std::env::var("REDIS_URL").is_err() {
|
|
return;
|
|
}
|
|
|
|
let config = FeatureStoreConfig {
|
|
postgres_url: None,
|
|
redis_url: Some(std::env::var("REDIS_URL").unwrap()),
|
|
vector_store_url: None,
|
|
enable_caching: true,
|
|
cache_ttl: Duration::from_secs(60),
|
|
max_connections: 10,
|
|
};
|
|
|
|
if let Ok(mut store) = FeatureStore::new(config).await {
|
|
// Test basic operations with Redis backend
|
|
store
|
|
.define_feature("redis_feat", serde_json::json!({"type": "string"}))
|
|
.await
|
|
.unwrap();
|
|
|
|
let feature = FeatureValue::String("redis_test".to_string());
|
|
store
|
|
.store_feature("redis_feat", "entity1", feature.clone())
|
|
.await
|
|
.unwrap();
|
|
|
|
let retrieved = store.get_feature("redis_feat", "entity1").await.unwrap();
|
|
assert_eq!(retrieved, feature);
|
|
|
|
// Clean up
|
|
store.delete_feature("redis_feat").await.unwrap();
|
|
}
|
|
}
|
|
}
|