use crate::config::StorageBackend;
use crate::database::{create_database, VirtualDatabase};
use crate::Result;
use mockforge_core::intelligent_behavior::session::SessionManager;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
pub struct SessionDataManager {
pub session_manager: Arc<SessionManager>,
#[allow(dead_code)]
storage_backend: StorageBackend,
session_databases: Arc<RwLock<HashMap<String, Arc<dyn VirtualDatabase + Send + Sync>>>>,
}
impl SessionDataManager {
pub fn new(session_manager: Arc<SessionManager>, storage_backend: StorageBackend) -> Self {
Self {
session_manager,
storage_backend,
session_databases: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn get_session_database(
&self,
session_id: &str,
) -> Result<Arc<dyn VirtualDatabase + Send + Sync>> {
let session_state = self.session_manager.get_session(session_id).await;
if session_state.is_none() {
return Err(crate::Error::internal(format!("Session '{}' not found", session_id)));
}
let databases = self.session_databases.read().await;
if let Some(db) = databases.get(session_id) {
return Ok(Arc::clone(db));
}
drop(databases);
let db_arc = create_database(&StorageBackend::Memory).await?;
let db_clone = Arc::clone(&db_arc);
let mut databases = self.session_databases.write().await;
databases.insert(session_id.to_string(), db_arc);
Ok(db_clone)
}
pub async fn cleanup_session_database(&self, session_id: &str) -> Result<()> {
let mut databases = self.session_databases.write().await;
databases.remove(session_id);
Ok(())
}
pub async fn cleanup_expired_sessions(&self) -> Result<usize> {
let _expired_count = self.session_manager.cleanup_expired_sessions().await;
let active_sessions = self.session_manager.list_sessions().await;
let active_set: std::collections::HashSet<String> = active_sessions.into_iter().collect();
let mut databases = self.session_databases.write().await;
let mut removed = 0;
let expired_db_sessions: Vec<String> = databases
.keys()
.filter(|session_id| !active_set.contains(*session_id))
.cloned()
.collect();
for session_id in expired_db_sessions {
databases.remove(&session_id);
removed += 1;
}
Ok(removed)
}
pub fn session_manager(&self) -> &Arc<SessionManager> {
&self.session_manager
}
pub async fn get_session_state(
&self,
session_id: &str,
) -> Option<mockforge_core::intelligent_behavior::types::SessionState> {
self.session_manager.get_session(session_id).await
}
pub async fn update_session_state(
&self,
session_id: &str,
state: mockforge_core::intelligent_behavior::types::SessionState,
) -> Result<()> {
self.session_manager.update_session(session_id, state).await
}
}
#[cfg(test)]
mod tests {
use super::*;
use mockforge_core::intelligent_behavior::session::SessionTracking;
fn create_test_session_manager() -> Arc<SessionManager> {
let config = SessionTracking::default();
Arc::new(SessionManager::new(config, 3600))
}
#[tokio::test]
async fn test_session_data_manager_new() {
let session_manager = create_test_session_manager();
let manager = SessionDataManager::new(session_manager.clone(), StorageBackend::Memory);
assert!(Arc::ptr_eq(&manager.session_manager, &session_manager));
}
#[tokio::test]
async fn test_get_session_database_no_session() {
let session_manager = create_test_session_manager();
let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
let result = manager.get_session_database("nonexistent-session").await;
assert!(result.is_err());
match result {
Err(err) => {
assert!(err.to_string().contains("Session"));
assert!(err.to_string().contains("not found"));
}
Ok(_) => panic!("Expected an error"),
}
}
#[tokio::test]
async fn test_get_session_database_creates_new() {
let session_manager = create_test_session_manager();
let _ = session_manager.get_or_create_session(Some("test-session".to_string())).await;
let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
let result = manager.get_session_database("test-session").await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_get_session_database_returns_same() {
let session_manager = create_test_session_manager();
let _ = session_manager.get_or_create_session(Some("test-session".to_string())).await;
let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
let db1 = manager.get_session_database("test-session").await.unwrap();
let db2 = manager.get_session_database("test-session").await.unwrap();
assert!(Arc::ptr_eq(&db1, &db2));
}
#[tokio::test]
async fn test_cleanup_session_database() {
let session_manager = create_test_session_manager();
let _ = session_manager.get_or_create_session(Some("test-session".to_string())).await;
let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
let _db = manager.get_session_database("test-session").await.unwrap();
let result = manager.cleanup_session_database("test-session").await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_cleanup_nonexistent_session() {
let session_manager = create_test_session_manager();
let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
let result = manager.cleanup_session_database("nonexistent").await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_session_manager_getter() {
let session_manager = create_test_session_manager();
let manager = SessionDataManager::new(session_manager.clone(), StorageBackend::Memory);
assert!(Arc::ptr_eq(manager.session_manager(), &session_manager));
}
#[tokio::test]
async fn test_get_session_state_exists() {
let session_manager = create_test_session_manager();
let _ = session_manager.get_or_create_session(Some("test-session".to_string())).await;
let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
let state = manager.get_session_state("test-session").await;
assert!(state.is_some());
}
#[tokio::test]
async fn test_get_session_state_not_exists() {
let session_manager = create_test_session_manager();
let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
let state = manager.get_session_state("nonexistent").await;
assert!(state.is_none());
}
#[tokio::test]
async fn test_multiple_sessions_isolated() {
let session_manager = create_test_session_manager();
let _ = session_manager.get_or_create_session(Some("session-1".to_string())).await;
let _ = session_manager.get_or_create_session(Some("session-2".to_string())).await;
let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
let db1 = manager.get_session_database("session-1").await.unwrap();
let db2 = manager.get_session_database("session-2").await.unwrap();
assert!(!Arc::ptr_eq(&db1, &db2));
}
#[tokio::test]
async fn test_cleanup_expired_sessions() {
let session_manager = create_test_session_manager();
let manager = SessionDataManager::new(session_manager, StorageBackend::Memory);
let result = manager.cleanup_expired_sessions().await;
assert!(result.is_ok());
}
}