graphmind/persistence/
tenant.rs

1//! Multi-tenancy implementation with resource quotas
2//!
3//! Implements REQ-TENANT-001 through REQ-TENANT-008
4
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7use std::sync::{Arc, RwLock};
8use thiserror::Error;
9// warn removed - was unused import causing compiler warning
10use tracing::{debug, info};
11
12/// Tenant errors
13#[derive(Error, Debug)]
14pub enum TenantError {
15    /// Tenant already exists
16    #[error("Tenant already exists: {0}")]
17    AlreadyExists(String),
18
19    /// Tenant not found
20    #[error("Tenant not found: {0}")]
21    NotFound(String),
22
23    /// Quota exceeded
24    #[error("Quota exceeded for tenant {tenant}: {resource}")]
25    QuotaExceeded { tenant: String, resource: String },
26
27    /// Permission denied
28    #[error("Permission denied for tenant {0}")]
29    PermissionDenied(String),
30}
31
32pub type TenantResult<T> = Result<T, TenantError>;
33
34/// Resource quotas for a tenant
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct ResourceQuotas {
37    /// Maximum number of nodes
38    pub max_nodes: Option<usize>,
39    /// Maximum number of edges
40    pub max_edges: Option<usize>,
41    /// Maximum memory in bytes
42    pub max_memory_bytes: Option<usize>,
43    /// Maximum storage in bytes
44    pub max_storage_bytes: Option<usize>,
45    /// Maximum concurrent connections
46    pub max_connections: Option<usize>,
47    /// Maximum query execution time in milliseconds
48    pub max_query_time_ms: Option<u64>,
49}
50
51impl Default for ResourceQuotas {
52    fn default() -> Self {
53        Self {
54            max_nodes: Some(1_000_000),              // 1M nodes
55            max_edges: Some(10_000_000),             // 10M edges
56            max_memory_bytes: Some(1_073_741_824),   // 1 GB
57            max_storage_bytes: Some(10_737_418_240), // 10 GB
58            max_connections: Some(100),
59            max_query_time_ms: Some(60_000), // 60 seconds
60        }
61    }
62}
63
64impl ResourceQuotas {
65    /// Create unlimited quotas
66    pub fn unlimited() -> Self {
67        Self {
68            max_nodes: None,
69            max_edges: None,
70            max_memory_bytes: None,
71            max_storage_bytes: None,
72            max_connections: None,
73            max_query_time_ms: None,
74        }
75    }
76}
77
78/// Current resource usage for a tenant
79#[derive(Debug, Clone, Default)]
80pub struct ResourceUsage {
81    /// Current number of nodes
82    pub node_count: usize,
83    /// Current number of edges
84    pub edge_count: usize,
85    /// Current memory usage in bytes
86    pub memory_bytes: usize,
87    /// Current storage usage in bytes
88    pub storage_bytes: usize,
89    /// Current number of connections
90    pub active_connections: usize,
91}
92
93impl ResourceUsage {
94    fn check_quota(&self, quotas: &ResourceQuotas, resource: &str) -> TenantResult<()> {
95        match resource {
96            "nodes" => {
97                if let Some(max) = quotas.max_nodes {
98                    if self.node_count >= max {
99                        return Err(TenantError::QuotaExceeded {
100                            tenant: String::new(),
101                            resource: format!("nodes ({}/{})", self.node_count, max),
102                        });
103                    }
104                }
105            }
106            "edges" => {
107                if let Some(max) = quotas.max_edges {
108                    if self.edge_count >= max {
109                        return Err(TenantError::QuotaExceeded {
110                            tenant: String::new(),
111                            resource: format!("edges ({}/{})", self.edge_count, max),
112                        });
113                    }
114                }
115            }
116            "memory" => {
117                if let Some(max) = quotas.max_memory_bytes {
118                    if self.memory_bytes >= max {
119                        return Err(TenantError::QuotaExceeded {
120                            tenant: String::new(),
121                            resource: format!("memory ({}/{})", self.memory_bytes, max),
122                        });
123                    }
124                }
125            }
126            "connections" => {
127                if let Some(max) = quotas.max_connections {
128                    if self.active_connections >= max {
129                        return Err(TenantError::QuotaExceeded {
130                            tenant: String::new(),
131                            resource: format!("connections ({}/{})", self.active_connections, max),
132                        });
133                    }
134                }
135            }
136            _ => {}
137        }
138        Ok(())
139    }
140}
141
142/// Tenant configuration and metadata
143#[derive(Debug, Clone, Serialize, Deserialize)]
144pub struct Tenant {
145    /// Tenant ID (unique identifier)
146    pub id: String,
147    /// Display name
148    pub name: String,
149    /// Creation timestamp
150    pub created_at: i64,
151    /// Resource quotas
152    pub quotas: ResourceQuotas,
153    /// Enabled status
154    pub enabled: bool,
155    /// Auto-Embed configuration
156    pub embed_config: Option<AutoEmbedConfig>,
157    /// NLQ configuration
158    pub nlq_config: Option<NLQConfig>,
159    /// Agent configuration
160    pub agent_config: Option<AgentConfig>,
161}
162
163impl Tenant {
164    /// Create a new tenant
165    pub fn new(id: String, name: String) -> Self {
166        Self {
167            id,
168            name,
169            created_at: chrono::Utc::now().timestamp(),
170            quotas: ResourceQuotas::default(),
171            enabled: true,
172            embed_config: None,
173            nlq_config: None,
174            agent_config: None,
175        }
176    }
177
178    /// Create a tenant with custom quotas
179    pub fn with_quotas(id: String, name: String, quotas: ResourceQuotas) -> Self {
180        Self {
181            id,
182            name,
183            created_at: chrono::Utc::now().timestamp(),
184            quotas,
185            enabled: true,
186            embed_config: None,
187            nlq_config: None,
188            agent_config: None,
189        }
190    }
191}
192
193/// LLM Provider options
194#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
195pub enum LLMProvider {
196    OpenAI,
197    Ollama,
198    Gemini,
199    AzureOpenAI,
200    Anthropic,
201    ClaudeCode,
202    Mock,
203}
204
205/// Tool definition for agents
206#[derive(Debug, Clone, Serialize, Deserialize)]
207pub struct ToolConfig {
208    pub name: String,
209    pub description: String,
210    pub parameters: serde_json::Value,
211    pub enabled: bool,
212}
213
214/// Configuration for Agentic features
215#[derive(Debug, Clone, Serialize, Deserialize)]
216pub struct AgentConfig {
217    /// Enabled status
218    pub enabled: bool,
219    /// LLM provider for the agent
220    pub provider: LLMProvider,
221    /// Model name
222    pub model: String,
223    /// API Key
224    pub api_key: Option<String>,
225    /// Base URL
226    pub api_base_url: Option<String>,
227    /// System prompt
228    pub system_prompt: Option<String>,
229    /// Available tools
230    pub tools: Vec<ToolConfig>,
231    /// Auto-trigger policies (e.g., on node creation)
232    pub policies: HashMap<String, String>, // Label -> Trigger Prompt
233}
234
235/// Configuration for NLQ features
236#[derive(Debug, Clone, Serialize, Deserialize)]
237pub struct NLQConfig {
238    /// Enabled status
239    pub enabled: bool,
240    /// The LLM provider to use
241    pub provider: LLMProvider,
242    /// Model name (e.g., "gpt-4o", "llama3")
243    pub model: String,
244    /// API Key (optional, can be loaded from env if None)
245    pub api_key: Option<String>,
246    /// API Base URL (required for Ollama/Azure, optional for others)
247    pub api_base_url: Option<String>,
248    /// System prompt for the LLM
249    pub system_prompt: Option<String>,
250}
251
252/// Configuration for Auto-Embed features
253#[derive(Debug, Clone, Serialize, Deserialize)]
254pub struct AutoEmbedConfig {
255    /// The LLM provider to use
256    pub provider: LLMProvider,
257    /// Model name (e.g., "text-embedding-3-small", "llama3")
258    pub embedding_model: String,
259    /// API Key (optional, can be loaded from env if None)
260    pub api_key: Option<String>,
261    /// API Base URL (required for Ollama/Azure, optional for others)
262    pub api_base_url: Option<String>,
263    /// Chunk size for text splitting
264    pub chunk_size: usize,
265    /// Overlap between chunks
266    pub chunk_overlap: usize,
267    /// Vector dimension size
268    pub vector_dimension: usize,
269    /// Embedding policies: Label -> `Vec<PropertyKey>`
270    pub embedding_policies: HashMap<String, Vec<String>>,
271}
272
273/// Tenant manager - manages all tenants and their resources
274pub struct TenantManager {
275    /// All tenants
276    tenants: Arc<RwLock<HashMap<String, Tenant>>>,
277    /// Resource usage per tenant
278    usage: Arc<RwLock<HashMap<String, ResourceUsage>>>,
279}
280
281impl TenantManager {
282    /// Create a new tenant manager
283    pub fn new() -> Self {
284        let mut tenants = HashMap::new();
285        let mut usage = HashMap::new();
286
287        // Create default tenant
288        let default_tenant = Tenant::new("default".to_string(), "Default Tenant".to_string());
289        tenants.insert("default".to_string(), default_tenant);
290        usage.insert("default".to_string(), ResourceUsage::default());
291
292        info!("Tenant manager initialized with default tenant");
293
294        Self {
295            tenants: Arc::new(RwLock::new(tenants)),
296            usage: Arc::new(RwLock::new(usage)),
297        }
298    }
299
300    /// Create a new tenant
301    pub fn create_tenant(
302        &self,
303        id: String,
304        name: String,
305        quotas: Option<ResourceQuotas>,
306    ) -> TenantResult<()> {
307        let mut tenants = self.tenants.write().unwrap();
308        let mut usage = self.usage.write().unwrap();
309
310        if tenants.contains_key(&id) {
311            return Err(TenantError::AlreadyExists(id));
312        }
313
314        let tenant = if let Some(quotas) = quotas {
315            Tenant::with_quotas(id.clone(), name, quotas)
316        } else {
317            Tenant::new(id.clone(), name)
318        };
319
320        tenants.insert(id.clone(), tenant);
321        usage.insert(id.clone(), ResourceUsage::default());
322
323        info!("Created tenant: {}", id);
324
325        Ok(())
326    }
327
328    /// Delete a tenant
329    pub fn delete_tenant(&self, id: &str) -> TenantResult<()> {
330        if id == "default" {
331            return Err(TenantError::PermissionDenied(
332                "Cannot delete default tenant".to_string(),
333            ));
334        }
335
336        let mut tenants = self.tenants.write().unwrap();
337        let mut usage = self.usage.write().unwrap();
338
339        if !tenants.contains_key(id) {
340            return Err(TenantError::NotFound(id.to_string()));
341        }
342
343        tenants.remove(id);
344        usage.remove(id);
345
346        info!("Deleted tenant: {}", id);
347
348        Ok(())
349    }
350
351    /// Get tenant information
352    pub fn get_tenant(&self, id: &str) -> TenantResult<Tenant> {
353        let tenants = self.tenants.read().unwrap();
354        tenants
355            .get(id)
356            .cloned()
357            .ok_or_else(|| TenantError::NotFound(id.to_string()))
358    }
359
360    /// List all tenants
361    pub fn list_tenants(&self) -> Vec<Tenant> {
362        let tenants = self.tenants.read().unwrap();
363        tenants.values().cloned().collect()
364    }
365
366    /// Check if a tenant exists and is enabled
367    pub fn is_tenant_enabled(&self, id: &str) -> bool {
368        let tenants = self.tenants.read().unwrap();
369        tenants.get(id).map(|t| t.enabled).unwrap_or(false)
370    }
371
372    /// Check and enforce resource quota
373    pub fn check_quota(&self, tenant_id: &str, resource: &str) -> TenantResult<()> {
374        let tenants = self.tenants.read().unwrap();
375        let usage = self.usage.read().unwrap();
376
377        let tenant = tenants
378            .get(tenant_id)
379            .ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
380
381        if !tenant.enabled {
382            return Err(TenantError::PermissionDenied(format!(
383                "Tenant {} is disabled",
384                tenant_id
385            )));
386        }
387
388        let current_usage = usage
389            .get(tenant_id)
390            .ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
391
392        current_usage
393            .check_quota(&tenant.quotas, resource)
394            .map_err(|e| match e {
395                TenantError::QuotaExceeded {
396                    tenant: _,
397                    resource,
398                } => TenantError::QuotaExceeded {
399                    tenant: tenant_id.to_string(),
400                    resource,
401                },
402                e => e,
403            })
404    }
405
406    /// Increment resource usage
407    pub fn increment_usage(
408        &self,
409        tenant_id: &str,
410        resource: &str,
411        amount: usize,
412    ) -> TenantResult<()> {
413        let mut usage = self.usage.write().unwrap();
414
415        let tenant_usage = usage
416            .get_mut(tenant_id)
417            .ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
418
419        match resource {
420            "nodes" => tenant_usage.node_count += amount,
421            "edges" => tenant_usage.edge_count += amount,
422            "memory" => tenant_usage.memory_bytes += amount,
423            "storage" => tenant_usage.storage_bytes += amount,
424            "connections" => tenant_usage.active_connections += amount,
425            _ => {}
426        }
427
428        debug!(
429            "Incremented {} for tenant {} by {}",
430            resource, tenant_id, amount
431        );
432
433        Ok(())
434    }
435
436    /// Decrement resource usage
437    pub fn decrement_usage(
438        &self,
439        tenant_id: &str,
440        resource: &str,
441        amount: usize,
442    ) -> TenantResult<()> {
443        let mut usage = self.usage.write().unwrap();
444
445        let tenant_usage = usage
446            .get_mut(tenant_id)
447            .ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
448
449        match resource {
450            "nodes" => tenant_usage.node_count = tenant_usage.node_count.saturating_sub(amount),
451            "edges" => tenant_usage.edge_count = tenant_usage.edge_count.saturating_sub(amount),
452            "memory" => {
453                tenant_usage.memory_bytes = tenant_usage.memory_bytes.saturating_sub(amount)
454            }
455            "storage" => {
456                tenant_usage.storage_bytes = tenant_usage.storage_bytes.saturating_sub(amount)
457            }
458            "connections" => {
459                tenant_usage.active_connections =
460                    tenant_usage.active_connections.saturating_sub(amount)
461            }
462            _ => {}
463        }
464
465        debug!(
466            "Decremented {} for tenant {} by {}",
467            resource, tenant_id, amount
468        );
469
470        Ok(())
471    }
472
473    /// Get resource usage for a tenant
474    pub fn get_usage(&self, tenant_id: &str) -> TenantResult<ResourceUsage> {
475        let usage = self.usage.read().unwrap();
476        usage
477            .get(tenant_id)
478            .cloned()
479            .ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))
480    }
481
482    /// Update tenant quotas
483    pub fn update_quotas(&self, tenant_id: &str, quotas: ResourceQuotas) -> TenantResult<()> {
484        let mut tenants = self.tenants.write().unwrap();
485
486        let tenant = tenants
487            .get_mut(tenant_id)
488            .ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
489
490        tenant.quotas = quotas;
491
492        info!("Updated quotas for tenant: {}", tenant_id);
493
494        Ok(())
495    }
496
497    /// Update Auto-Embed configuration for a tenant
498    pub fn update_embed_config(
499        &self,
500        tenant_id: &str,
501        config: Option<AutoEmbedConfig>,
502    ) -> TenantResult<()> {
503        let mut tenants = self.tenants.write().unwrap();
504
505        let tenant = tenants
506            .get_mut(tenant_id)
507            .ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
508
509        tenant.embed_config = config;
510
511        info!("Updated Auto-Embed config for tenant: {}", tenant_id);
512
513        Ok(())
514    }
515
516    /// Update NLQ configuration for a tenant
517    pub fn update_nlq_config(
518        &self,
519        tenant_id: &str,
520        config: Option<NLQConfig>,
521    ) -> TenantResult<()> {
522        let mut tenants = self.tenants.write().unwrap();
523
524        let tenant = tenants
525            .get_mut(tenant_id)
526            .ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
527
528        tenant.nlq_config = config;
529
530        info!("Updated NLQ config for tenant: {}", tenant_id);
531
532        Ok(())
533    }
534
535    /// Update Agent configuration for a tenant
536    pub fn update_agent_config(
537        &self,
538        tenant_id: &str,
539        config: Option<AgentConfig>,
540    ) -> TenantResult<()> {
541        let mut tenants = self.tenants.write().unwrap();
542
543        let tenant = tenants
544            .get_mut(tenant_id)
545            .ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
546
547        tenant.agent_config = config;
548
549        info!("Updated Agent config for tenant: {}", tenant_id);
550
551        Ok(())
552    }
553
554    /// Enable/disable a tenant
555    pub fn set_enabled(&self, tenant_id: &str, enabled: bool) -> TenantResult<()> {
556        if tenant_id == "default" {
557            return Err(TenantError::PermissionDenied(
558                "Cannot disable default tenant".to_string(),
559            ));
560        }
561
562        let mut tenants = self.tenants.write().unwrap();
563
564        let tenant = tenants
565            .get_mut(tenant_id)
566            .ok_or_else(|| TenantError::NotFound(tenant_id.to_string()))?;
567
568        tenant.enabled = enabled;
569
570        info!("Set tenant {} enabled status to: {}", tenant_id, enabled);
571
572        Ok(())
573    }
574}
575
576impl Default for TenantManager {
577    fn default() -> Self {
578        Self::new()
579    }
580}
581
582#[cfg(test)]
583mod tests {
584    use super::*;
585
586    #[test]
587    fn test_tenant_manager_creation() {
588        let manager = TenantManager::new();
589        assert!(manager.is_tenant_enabled("default"));
590    }
591
592    #[test]
593    fn test_create_tenant() {
594        let manager = TenantManager::new();
595        manager
596            .create_tenant("tenant1".to_string(), "Tenant 1".to_string(), None)
597            .unwrap();
598
599        let tenant = manager.get_tenant("tenant1").unwrap();
600        assert_eq!(tenant.id, "tenant1");
601        assert_eq!(tenant.name, "Tenant 1");
602        assert!(tenant.enabled);
603    }
604
605    #[test]
606    fn test_duplicate_tenant() {
607        let manager = TenantManager::new();
608        manager
609            .create_tenant("tenant1".to_string(), "Tenant 1".to_string(), None)
610            .unwrap();
611
612        let result = manager.create_tenant(
613            "tenant1".to_string(),
614            "Tenant 1 Duplicate".to_string(),
615            None,
616        );
617        assert!(result.is_err());
618        assert!(matches!(result.unwrap_err(), TenantError::AlreadyExists(_)));
619    }
620
621    #[test]
622    fn test_delete_tenant() {
623        let manager = TenantManager::new();
624        manager
625            .create_tenant("tenant1".to_string(), "Tenant 1".to_string(), None)
626            .unwrap();
627
628        manager.delete_tenant("tenant1").unwrap();
629        assert!(manager.get_tenant("tenant1").is_err());
630    }
631
632    #[test]
633    fn test_cannot_delete_default() {
634        let manager = TenantManager::new();
635        let result = manager.delete_tenant("default");
636        assert!(result.is_err());
637    }
638
639    #[test]
640    fn test_quota_enforcement() {
641        let manager = TenantManager::new();
642
643        let mut quotas = ResourceQuotas::default();
644        quotas.max_nodes = Some(10);
645
646        manager
647            .create_tenant("tenant1".to_string(), "Tenant 1".to_string(), Some(quotas))
648            .unwrap();
649
650        // Should succeed for first 10 nodes
651        for _ in 0..10 {
652            manager.check_quota("tenant1", "nodes").unwrap();
653            manager.increment_usage("tenant1", "nodes", 1).unwrap();
654        }
655
656        // 11th should fail
657        let result = manager.check_quota("tenant1", "nodes");
658        assert!(result.is_err());
659        assert!(matches!(
660            result.unwrap_err(),
661            TenantError::QuotaExceeded { .. }
662        ));
663    }
664
665    #[test]
666    fn test_usage_tracking() {
667        let manager = TenantManager::new();
668        manager
669            .create_tenant("tenant1".to_string(), "Tenant 1".to_string(), None)
670            .unwrap();
671
672        manager.increment_usage("tenant1", "nodes", 5).unwrap();
673        manager.increment_usage("tenant1", "edges", 10).unwrap();
674
675        let usage = manager.get_usage("tenant1").unwrap();
676        assert_eq!(usage.node_count, 5);
677        assert_eq!(usage.edge_count, 10);
678
679        manager.decrement_usage("tenant1", "nodes", 2).unwrap();
680        let usage = manager.get_usage("tenant1").unwrap();
681        assert_eq!(usage.node_count, 3);
682    }
683
684    #[test]
685    fn test_list_tenants() {
686        let manager = TenantManager::new();
687        manager
688            .create_tenant("tenant1".to_string(), "Tenant 1".to_string(), None)
689            .unwrap();
690        manager
691            .create_tenant("tenant2".to_string(), "Tenant 2".to_string(), None)
692            .unwrap();
693
694        let tenants = manager.list_tenants();
695        assert_eq!(tenants.len(), 3); // default + 2 new
696    }
697
698    #[test]
699    fn test_disable_tenant() {
700        let manager = TenantManager::new();
701        manager
702            .create_tenant("tenant1".to_string(), "Tenant 1".to_string(), None)
703            .unwrap();
704
705        manager.set_enabled("tenant1", false).unwrap();
706        assert!(!manager.is_tenant_enabled("tenant1"));
707
708        let result = manager.check_quota("tenant1", "nodes");
709        assert!(result.is_err());
710        assert!(matches!(
711            result.unwrap_err(),
712            TenantError::PermissionDenied(_)
713        ));
714    }
715
716    #[test]
717    fn test_update_embed_config() {
718        let manager = TenantManager::new();
719        manager
720            .create_tenant("tenant1".to_string(), "Tenant 1".to_string(), None)
721            .unwrap();
722
723        let embed_config = AutoEmbedConfig {
724            provider: LLMProvider::OpenAI,
725            embedding_model: "text-embedding-3-small".to_string(),
726            api_key: Some("sk-test".to_string()),
727            api_base_url: None,
728            chunk_size: 512,
729            chunk_overlap: 64,
730            vector_dimension: 1536,
731            embedding_policies: HashMap::from([(
732                "Document".to_string(),
733                vec!["content".to_string()],
734            )]),
735        };
736
737        manager
738            .update_embed_config("tenant1", Some(embed_config))
739            .unwrap();
740
741        let tenant = manager.get_tenant("tenant1").unwrap();
742        assert!(tenant.embed_config.is_some());
743        let config = tenant.embed_config.unwrap();
744        assert_eq!(config.provider, LLMProvider::OpenAI);
745        assert_eq!(config.embedding_model, "text-embedding-3-small");
746    }
747
748    // ========== Batch 7: Additional Tenant Tests ==========
749
750    #[test]
751    fn test_update_nlq_config() {
752        let manager = TenantManager::new();
753        manager
754            .create_tenant("tenant1".to_string(), "T1".to_string(), None)
755            .unwrap();
756
757        let nlq_config = NLQConfig {
758            enabled: true,
759            provider: LLMProvider::Ollama,
760            model: "llama3".to_string(),
761            api_key: None,
762            api_base_url: Some("http://localhost:11434".to_string()),
763            system_prompt: Some("You are a Cypher expert.".to_string()),
764        };
765
766        manager
767            .update_nlq_config("tenant1", Some(nlq_config))
768            .unwrap();
769
770        let tenant = manager.get_tenant("tenant1").unwrap();
771        assert!(tenant.nlq_config.is_some());
772        let config = tenant.nlq_config.unwrap();
773        assert_eq!(config.provider, LLMProvider::Ollama);
774        assert_eq!(config.model, "llama3");
775    }
776
777    #[test]
778    fn test_update_agent_config() {
779        let manager = TenantManager::new();
780        manager
781            .create_tenant("tenant1".to_string(), "T1".to_string(), None)
782            .unwrap();
783
784        let agent_config = AgentConfig {
785            enabled: true,
786            provider: LLMProvider::Mock,
787            model: "mock".to_string(),
788            api_key: None,
789            api_base_url: None,
790            system_prompt: None,
791            tools: vec![],
792            policies: HashMap::new(),
793        };
794
795        manager
796            .update_agent_config("tenant1", Some(agent_config))
797            .unwrap();
798
799        let tenant = manager.get_tenant("tenant1").unwrap();
800        assert!(tenant.agent_config.is_some());
801        assert_eq!(tenant.agent_config.unwrap().provider, LLMProvider::Mock);
802    }
803
804    #[test]
805    fn test_update_config_nonexistent_tenant() {
806        let manager = TenantManager::new();
807
808        let nlq_config = NLQConfig {
809            enabled: true,
810            provider: LLMProvider::Mock,
811            model: "mock".to_string(),
812            api_key: None,
813            api_base_url: None,
814            system_prompt: None,
815        };
816
817        let result = manager.update_nlq_config("nonexistent", Some(nlq_config));
818        assert!(result.is_err());
819    }
820
821    // ========== Additional Tenant Coverage Tests ==========
822
823    #[test]
824    fn test_resource_quotas_default() {
825        let q = ResourceQuotas::default();
826        assert_eq!(q.max_nodes, Some(1_000_000));
827        assert_eq!(q.max_edges, Some(10_000_000));
828        assert_eq!(q.max_memory_bytes, Some(1_073_741_824));
829        assert_eq!(q.max_storage_bytes, Some(10_737_418_240));
830        assert_eq!(q.max_connections, Some(100));
831        assert_eq!(q.max_query_time_ms, Some(60_000));
832    }
833
834    #[test]
835    fn test_resource_quotas_unlimited() {
836        let q = ResourceQuotas::unlimited();
837        assert!(q.max_nodes.is_none());
838        assert!(q.max_edges.is_none());
839        assert!(q.max_memory_bytes.is_none());
840        assert!(q.max_storage_bytes.is_none());
841        assert!(q.max_connections.is_none());
842        assert!(q.max_query_time_ms.is_none());
843    }
844
845    #[test]
846    fn test_tenant_new() {
847        let t = Tenant::new("test_id".to_string(), "Test Name".to_string());
848        assert_eq!(t.id, "test_id");
849        assert_eq!(t.name, "Test Name");
850        assert!(t.enabled);
851        assert!(t.embed_config.is_none());
852        assert!(t.nlq_config.is_none());
853        assert!(t.agent_config.is_none());
854        assert!(t.created_at > 0);
855    }
856
857    #[test]
858    fn test_tenant_with_quotas() {
859        let quotas = ResourceQuotas::unlimited();
860        let t = Tenant::with_quotas("custom".to_string(), "Custom".to_string(), quotas);
861        assert_eq!(t.id, "custom");
862        assert!(t.quotas.max_nodes.is_none());
863        assert!(t.enabled);
864    }
865
866    #[test]
867    fn test_tenant_manager_default() {
868        let manager = TenantManager::default();
869        assert!(manager.is_tenant_enabled("default"));
870    }
871
872    #[test]
873    fn test_delete_nonexistent_tenant() {
874        let manager = TenantManager::new();
875        let result = manager.delete_tenant("ghost");
876        assert!(result.is_err());
877        assert!(matches!(result.unwrap_err(), TenantError::NotFound(_)));
878    }
879
880    #[test]
881    fn test_get_nonexistent_tenant() {
882        let manager = TenantManager::new();
883        let result = manager.get_tenant("ghost");
884        assert!(result.is_err());
885        assert!(matches!(result.unwrap_err(), TenantError::NotFound(_)));
886    }
887
888    #[test]
889    fn test_is_tenant_enabled_nonexistent() {
890        let manager = TenantManager::new();
891        assert!(!manager.is_tenant_enabled("ghost"));
892    }
893
894    #[test]
895    fn test_cannot_disable_default_tenant() {
896        let manager = TenantManager::new();
897        let result = manager.set_enabled("default", false);
898        assert!(result.is_err());
899        assert!(matches!(
900            result.unwrap_err(),
901            TenantError::PermissionDenied(_)
902        ));
903    }
904
905    #[test]
906    fn test_set_enabled_nonexistent() {
907        let manager = TenantManager::new();
908        let result = manager.set_enabled("ghost", true);
909        assert!(result.is_err());
910        assert!(matches!(result.unwrap_err(), TenantError::NotFound(_)));
911    }
912
913    #[test]
914    fn test_set_enabled_reenable() {
915        let manager = TenantManager::new();
916        manager
917            .create_tenant("t1".to_string(), "T1".to_string(), None)
918            .unwrap();
919
920        manager.set_enabled("t1", false).unwrap();
921        assert!(!manager.is_tenant_enabled("t1"));
922
923        manager.set_enabled("t1", true).unwrap();
924        assert!(manager.is_tenant_enabled("t1"));
925    }
926
927    #[test]
928    fn test_update_quotas() {
929        let manager = TenantManager::new();
930        manager
931            .create_tenant("t1".to_string(), "T1".to_string(), None)
932            .unwrap();
933
934        let new_quotas = ResourceQuotas {
935            max_nodes: Some(500),
936            max_edges: Some(1000),
937            max_memory_bytes: None,
938            max_storage_bytes: None,
939            max_connections: Some(10),
940            max_query_time_ms: Some(5000),
941        };
942
943        manager.update_quotas("t1", new_quotas).unwrap();
944
945        let tenant = manager.get_tenant("t1").unwrap();
946        assert_eq!(tenant.quotas.max_nodes, Some(500));
947        assert_eq!(tenant.quotas.max_edges, Some(1000));
948        assert!(tenant.quotas.max_memory_bytes.is_none());
949    }
950
951    #[test]
952    fn test_update_quotas_nonexistent() {
953        let manager = TenantManager::new();
954        let result = manager.update_quotas("ghost", ResourceQuotas::unlimited());
955        assert!(result.is_err());
956    }
957
958    #[test]
959    fn test_check_quota_disabled_tenant() {
960        let manager = TenantManager::new();
961        manager
962            .create_tenant("t1".to_string(), "T1".to_string(), None)
963            .unwrap();
964        manager.set_enabled("t1", false).unwrap();
965
966        let result = manager.check_quota("t1", "nodes");
967        assert!(result.is_err());
968        assert!(matches!(
969            result.unwrap_err(),
970            TenantError::PermissionDenied(_)
971        ));
972    }
973
974    #[test]
975    fn test_check_quota_nonexistent_tenant() {
976        let manager = TenantManager::new();
977        let result = manager.check_quota("ghost", "nodes");
978        assert!(result.is_err());
979    }
980
981    #[test]
982    fn test_quota_enforcement_edges() {
983        let manager = TenantManager::new();
984        let quotas = ResourceQuotas {
985            max_nodes: None,
986            max_edges: Some(5),
987            max_memory_bytes: None,
988            max_storage_bytes: None,
989            max_connections: None,
990            max_query_time_ms: None,
991        };
992        manager
993            .create_tenant("t1".to_string(), "T1".to_string(), Some(quotas))
994            .unwrap();
995
996        for _ in 0..5 {
997            manager.check_quota("t1", "edges").unwrap();
998            manager.increment_usage("t1", "edges", 1).unwrap();
999        }
1000
1001        let result = manager.check_quota("t1", "edges");
1002        assert!(result.is_err());
1003    }
1004
1005    #[test]
1006    fn test_quota_enforcement_memory() {
1007        let manager = TenantManager::new();
1008        let quotas = ResourceQuotas {
1009            max_nodes: None,
1010            max_edges: None,
1011            max_memory_bytes: Some(1024),
1012            max_storage_bytes: None,
1013            max_connections: None,
1014            max_query_time_ms: None,
1015        };
1016        manager
1017            .create_tenant("t1".to_string(), "T1".to_string(), Some(quotas))
1018            .unwrap();
1019
1020        manager.increment_usage("t1", "memory", 1024).unwrap();
1021        let result = manager.check_quota("t1", "memory");
1022        assert!(result.is_err());
1023    }
1024
1025    #[test]
1026    fn test_quota_enforcement_connections() {
1027        let manager = TenantManager::new();
1028        let quotas = ResourceQuotas {
1029            max_nodes: None,
1030            max_edges: None,
1031            max_memory_bytes: None,
1032            max_storage_bytes: None,
1033            max_connections: Some(2),
1034            max_query_time_ms: None,
1035        };
1036        manager
1037            .create_tenant("t1".to_string(), "T1".to_string(), Some(quotas))
1038            .unwrap();
1039
1040        manager.increment_usage("t1", "connections", 2).unwrap();
1041        let result = manager.check_quota("t1", "connections");
1042        assert!(result.is_err());
1043    }
1044
1045    #[test]
1046    fn test_quota_unlimited_allows_everything() {
1047        let manager = TenantManager::new();
1048        manager
1049            .create_tenant(
1050                "t1".to_string(),
1051                "T1".to_string(),
1052                Some(ResourceQuotas::unlimited()),
1053            )
1054            .unwrap();
1055
1056        manager.increment_usage("t1", "nodes", 999_999).unwrap();
1057        assert!(manager.check_quota("t1", "nodes").is_ok());
1058
1059        manager.increment_usage("t1", "edges", 999_999).unwrap();
1060        assert!(manager.check_quota("t1", "edges").is_ok());
1061    }
1062
1063    #[test]
1064    fn test_increment_usage_unknown_resource() {
1065        let manager = TenantManager::new();
1066        manager
1067            .create_tenant("t1".to_string(), "T1".to_string(), None)
1068            .unwrap();
1069        // Unknown resource should not panic, just be a no-op
1070        manager
1071            .increment_usage("t1", "unknown_resource", 100)
1072            .unwrap();
1073    }
1074
1075    #[test]
1076    fn test_decrement_usage_unknown_resource() {
1077        let manager = TenantManager::new();
1078        manager
1079            .create_tenant("t1".to_string(), "T1".to_string(), None)
1080            .unwrap();
1081        manager
1082            .decrement_usage("t1", "unknown_resource", 100)
1083            .unwrap();
1084    }
1085
1086    #[test]
1087    fn test_increment_usage_nonexistent_tenant() {
1088        let manager = TenantManager::new();
1089        let result = manager.increment_usage("ghost", "nodes", 1);
1090        assert!(result.is_err());
1091    }
1092
1093    #[test]
1094    fn test_decrement_usage_nonexistent_tenant() {
1095        let manager = TenantManager::new();
1096        let result = manager.decrement_usage("ghost", "nodes", 1);
1097        assert!(result.is_err());
1098    }
1099
1100    #[test]
1101    fn test_decrement_usage_saturating() {
1102        let manager = TenantManager::new();
1103        manager
1104            .create_tenant("t1".to_string(), "T1".to_string(), None)
1105            .unwrap();
1106
1107        manager.increment_usage("t1", "nodes", 3).unwrap();
1108        // Decrement more than the current count: should saturate to 0
1109        manager.decrement_usage("t1", "nodes", 100).unwrap();
1110
1111        let usage = manager.get_usage("t1").unwrap();
1112        assert_eq!(usage.node_count, 0);
1113    }
1114
1115    #[test]
1116    fn test_decrement_all_resource_types() {
1117        let manager = TenantManager::new();
1118        manager
1119            .create_tenant("t1".to_string(), "T1".to_string(), None)
1120            .unwrap();
1121
1122        manager.increment_usage("t1", "nodes", 10).unwrap();
1123        manager.increment_usage("t1", "edges", 20).unwrap();
1124        manager.increment_usage("t1", "memory", 1000).unwrap();
1125        manager.increment_usage("t1", "storage", 2000).unwrap();
1126        manager.increment_usage("t1", "connections", 5).unwrap();
1127
1128        manager.decrement_usage("t1", "nodes", 3).unwrap();
1129        manager.decrement_usage("t1", "edges", 5).unwrap();
1130        manager.decrement_usage("t1", "memory", 200).unwrap();
1131        manager.decrement_usage("t1", "storage", 500).unwrap();
1132        manager.decrement_usage("t1", "connections", 2).unwrap();
1133
1134        let usage = manager.get_usage("t1").unwrap();
1135        assert_eq!(usage.node_count, 7);
1136        assert_eq!(usage.edge_count, 15);
1137        assert_eq!(usage.memory_bytes, 800);
1138        assert_eq!(usage.storage_bytes, 1500);
1139        assert_eq!(usage.active_connections, 3);
1140    }
1141
1142    #[test]
1143    fn test_get_usage_nonexistent_tenant() {
1144        let manager = TenantManager::new();
1145        let result = manager.get_usage("ghost");
1146        assert!(result.is_err());
1147    }
1148
1149    #[test]
1150    fn test_get_usage_default_tenant() {
1151        let manager = TenantManager::new();
1152        let usage = manager.get_usage("default").unwrap();
1153        assert_eq!(usage.node_count, 0);
1154        assert_eq!(usage.edge_count, 0);
1155        assert_eq!(usage.memory_bytes, 0);
1156        assert_eq!(usage.storage_bytes, 0);
1157        assert_eq!(usage.active_connections, 0);
1158    }
1159
1160    #[test]
1161    fn test_update_embed_config_nonexistent() {
1162        let manager = TenantManager::new();
1163        let config = AutoEmbedConfig {
1164            provider: LLMProvider::OpenAI,
1165            embedding_model: "text-embedding-3-small".to_string(),
1166            api_key: None,
1167            api_base_url: None,
1168            chunk_size: 256,
1169            chunk_overlap: 32,
1170            vector_dimension: 1536,
1171            embedding_policies: HashMap::new(),
1172        };
1173        let result = manager.update_embed_config("ghost", Some(config));
1174        assert!(result.is_err());
1175    }
1176
1177    #[test]
1178    fn test_update_agent_config_nonexistent() {
1179        let manager = TenantManager::new();
1180        let result = manager.update_agent_config("ghost", None);
1181        assert!(result.is_err());
1182    }
1183
1184    #[test]
1185    fn test_clear_embed_config() {
1186        let manager = TenantManager::new();
1187        manager
1188            .create_tenant("t1".to_string(), "T1".to_string(), None)
1189            .unwrap();
1190
1191        let config = AutoEmbedConfig {
1192            provider: LLMProvider::OpenAI,
1193            embedding_model: "test".to_string(),
1194            api_key: None,
1195            api_base_url: None,
1196            chunk_size: 256,
1197            chunk_overlap: 32,
1198            vector_dimension: 768,
1199            embedding_policies: HashMap::new(),
1200        };
1201        manager.update_embed_config("t1", Some(config)).unwrap();
1202        assert!(manager.get_tenant("t1").unwrap().embed_config.is_some());
1203
1204        // Clear it
1205        manager.update_embed_config("t1", None).unwrap();
1206        assert!(manager.get_tenant("t1").unwrap().embed_config.is_none());
1207    }
1208
1209    #[test]
1210    fn test_clear_nlq_config() {
1211        let manager = TenantManager::new();
1212        manager
1213            .create_tenant("t1".to_string(), "T1".to_string(), None)
1214            .unwrap();
1215
1216        let config = NLQConfig {
1217            enabled: true,
1218            provider: LLMProvider::Mock,
1219            model: "mock".to_string(),
1220            api_key: None,
1221            api_base_url: None,
1222            system_prompt: None,
1223        };
1224        manager.update_nlq_config("t1", Some(config)).unwrap();
1225        assert!(manager.get_tenant("t1").unwrap().nlq_config.is_some());
1226
1227        manager.update_nlq_config("t1", None).unwrap();
1228        assert!(manager.get_tenant("t1").unwrap().nlq_config.is_none());
1229    }
1230
1231    #[test]
1232    fn test_clear_agent_config() {
1233        let manager = TenantManager::new();
1234        manager
1235            .create_tenant("t1".to_string(), "T1".to_string(), None)
1236            .unwrap();
1237
1238        let config = AgentConfig {
1239            enabled: true,
1240            provider: LLMProvider::Mock,
1241            model: "mock".to_string(),
1242            api_key: None,
1243            api_base_url: None,
1244            system_prompt: None,
1245            tools: vec![],
1246            policies: HashMap::new(),
1247        };
1248        manager.update_agent_config("t1", Some(config)).unwrap();
1249        assert!(manager.get_tenant("t1").unwrap().agent_config.is_some());
1250
1251        manager.update_agent_config("t1", None).unwrap();
1252        assert!(manager.get_tenant("t1").unwrap().agent_config.is_none());
1253    }
1254
1255    #[test]
1256    fn test_llm_provider_variants() {
1257        // Ensure all LLMProvider variants are distinct
1258        let providers = vec![
1259            LLMProvider::OpenAI,
1260            LLMProvider::Ollama,
1261            LLMProvider::Gemini,
1262            LLMProvider::AzureOpenAI,
1263            LLMProvider::Anthropic,
1264            LLMProvider::ClaudeCode,
1265            LLMProvider::Mock,
1266        ];
1267        for (i, p1) in providers.iter().enumerate() {
1268            for (j, p2) in providers.iter().enumerate() {
1269                if i == j {
1270                    assert_eq!(p1, p2);
1271                } else {
1272                    assert_ne!(p1, p2);
1273                }
1274            }
1275        }
1276    }
1277
1278    #[test]
1279    fn test_resource_quotas_serialization() {
1280        let quotas = ResourceQuotas {
1281            max_nodes: Some(500),
1282            max_edges: Some(1000),
1283            max_memory_bytes: None,
1284            max_storage_bytes: None,
1285            max_connections: Some(10),
1286            max_query_time_ms: Some(30_000),
1287        };
1288        let json = serde_json::to_string(&quotas).unwrap();
1289        let deserialized: ResourceQuotas = serde_json::from_str(&json).unwrap();
1290        assert_eq!(deserialized.max_nodes, Some(500));
1291        assert_eq!(deserialized.max_edges, Some(1000));
1292        assert!(deserialized.max_memory_bytes.is_none());
1293        assert_eq!(deserialized.max_connections, Some(10));
1294    }
1295
1296    #[test]
1297    fn test_nlq_config_serialization() {
1298        let config = NLQConfig {
1299            enabled: true,
1300            provider: LLMProvider::Gemini,
1301            model: "gemini-pro".to_string(),
1302            api_key: Some("key123".to_string()),
1303            api_base_url: None,
1304            system_prompt: Some("You are a graph expert.".to_string()),
1305        };
1306        let json = serde_json::to_string(&config).unwrap();
1307        let deserialized: NLQConfig = serde_json::from_str(&json).unwrap();
1308        assert_eq!(deserialized.provider, LLMProvider::Gemini);
1309        assert_eq!(deserialized.model, "gemini-pro");
1310        assert_eq!(deserialized.api_key, Some("key123".to_string()));
1311    }
1312
1313    #[test]
1314    fn test_auto_embed_config_serialization() {
1315        let config = AutoEmbedConfig {
1316            provider: LLMProvider::OpenAI,
1317            embedding_model: "text-embedding-3-small".to_string(),
1318            api_key: Some("sk-test".to_string()),
1319            api_base_url: None,
1320            chunk_size: 512,
1321            chunk_overlap: 64,
1322            vector_dimension: 1536,
1323            embedding_policies: HashMap::from([(
1324                "Document".to_string(),
1325                vec!["content".to_string(), "title".to_string()],
1326            )]),
1327        };
1328        let json = serde_json::to_string(&config).unwrap();
1329        let deserialized: AutoEmbedConfig = serde_json::from_str(&json).unwrap();
1330        assert_eq!(deserialized.chunk_size, 512);
1331        assert_eq!(deserialized.vector_dimension, 1536);
1332        assert_eq!(deserialized.embedding_policies.len(), 1);
1333    }
1334
1335    #[test]
1336    fn test_agent_config_serialization() {
1337        let config = AgentConfig {
1338            enabled: true,
1339            provider: LLMProvider::Anthropic,
1340            model: "claude-3-opus".to_string(),
1341            api_key: Some("sk-ant-test".to_string()),
1342            api_base_url: None,
1343            system_prompt: Some("You are an agent.".to_string()),
1344            tools: vec![ToolConfig {
1345                name: "search".to_string(),
1346                description: "Search the graph".to_string(),
1347                parameters: serde_json::json!({"query": "string"}),
1348                enabled: true,
1349            }],
1350            policies: HashMap::from([("Person".to_string(), "Enrich person data".to_string())]),
1351        };
1352        let json = serde_json::to_string(&config).unwrap();
1353        let deserialized: AgentConfig = serde_json::from_str(&json).unwrap();
1354        assert_eq!(deserialized.provider, LLMProvider::Anthropic);
1355        assert_eq!(deserialized.tools.len(), 1);
1356        assert_eq!(deserialized.tools[0].name, "search");
1357        assert!(deserialized.tools[0].enabled);
1358        assert_eq!(deserialized.policies.len(), 1);
1359    }
1360
1361    #[test]
1362    fn test_tool_config_serialization() {
1363        let tool = ToolConfig {
1364            name: "query".to_string(),
1365            description: "Execute a Cypher query".to_string(),
1366            parameters: serde_json::json!({"cypher": "string", "readonly": "boolean"}),
1367            enabled: false,
1368        };
1369        let json = serde_json::to_string(&tool).unwrap();
1370        let deserialized: ToolConfig = serde_json::from_str(&json).unwrap();
1371        assert_eq!(deserialized.name, "query");
1372        assert!(!deserialized.enabled);
1373    }
1374
1375    #[test]
1376    fn test_tenant_serialization() {
1377        let tenant = Tenant::new("t1".to_string(), "Tenant 1".to_string());
1378        let json = serde_json::to_string(&tenant).unwrap();
1379        let deserialized: Tenant = serde_json::from_str(&json).unwrap();
1380        assert_eq!(deserialized.id, "t1");
1381        assert_eq!(deserialized.name, "Tenant 1");
1382        assert!(deserialized.enabled);
1383    }
1384
1385    #[test]
1386    fn test_check_quota_unknown_resource_passes() {
1387        let manager = TenantManager::new();
1388        // Unknown resource type should pass quota check (no limits defined)
1389        let result = manager.check_quota("default", "something_random");
1390        assert!(result.is_ok());
1391    }
1392
1393    #[test]
1394    fn test_increment_storage_usage() {
1395        let manager = TenantManager::new();
1396        manager
1397            .create_tenant("t1".to_string(), "T1".to_string(), None)
1398            .unwrap();
1399
1400        manager.increment_usage("t1", "storage", 5000).unwrap();
1401        let usage = manager.get_usage("t1").unwrap();
1402        assert_eq!(usage.storage_bytes, 5000);
1403    }
1404
1405    #[test]
1406    fn test_resource_usage_default() {
1407        let usage = ResourceUsage::default();
1408        assert_eq!(usage.node_count, 0);
1409        assert_eq!(usage.edge_count, 0);
1410        assert_eq!(usage.memory_bytes, 0);
1411        assert_eq!(usage.storage_bytes, 0);
1412        assert_eq!(usage.active_connections, 0);
1413    }
1414
1415    #[test]
1416    fn test_llm_provider_serialization_roundtrip() {
1417        let providers = vec![
1418            LLMProvider::OpenAI,
1419            LLMProvider::Ollama,
1420            LLMProvider::Gemini,
1421            LLMProvider::AzureOpenAI,
1422            LLMProvider::Anthropic,
1423            LLMProvider::ClaudeCode,
1424            LLMProvider::Mock,
1425        ];
1426        for provider in providers {
1427            let json = serde_json::to_string(&provider).unwrap();
1428            let deserialized: LLMProvider = serde_json::from_str(&json).unwrap();
1429            assert_eq!(provider, deserialized);
1430        }
1431    }
1432
1433    // ========== Additional Tenant Coverage Tests ==========
1434
1435    #[test]
1436    fn test_tenant_error_display_messages() {
1437        let e1 = TenantError::AlreadyExists("t1".to_string());
1438        assert!(format!("{}", e1).contains("already exists"));
1439        assert!(format!("{}", e1).contains("t1"));
1440
1441        let e2 = TenantError::NotFound("t2".to_string());
1442        assert!(format!("{}", e2).contains("not found"));
1443        assert!(format!("{}", e2).contains("t2"));
1444
1445        let e3 = TenantError::QuotaExceeded {
1446            tenant: "t3".to_string(),
1447            resource: "nodes (100/100)".to_string(),
1448        };
1449        assert!(format!("{}", e3).contains("Quota exceeded"));
1450        assert!(format!("{}", e3).contains("t3"));
1451
1452        let e4 = TenantError::PermissionDenied("t4".to_string());
1453        assert!(format!("{}", e4).contains("Permission denied"));
1454    }
1455
1456    #[test]
1457    fn test_auto_embed_config_construction() {
1458        let mut policies = HashMap::new();
1459        policies.insert(
1460            "Article".to_string(),
1461            vec!["title".to_string(), "body".to_string()],
1462        );
1463        policies.insert("Comment".to_string(), vec!["text".to_string()]);
1464
1465        let config = AutoEmbedConfig {
1466            provider: LLMProvider::Ollama,
1467            embedding_model: "nomic-embed-text".to_string(),
1468            api_key: None,
1469            api_base_url: Some("http://localhost:11434".to_string()),
1470            chunk_size: 1024,
1471            chunk_overlap: 128,
1472            vector_dimension: 768,
1473            embedding_policies: policies,
1474        };
1475
1476        assert_eq!(config.provider, LLMProvider::Ollama);
1477        assert_eq!(config.chunk_size, 1024);
1478        assert_eq!(config.chunk_overlap, 128);
1479        assert_eq!(config.vector_dimension, 768);
1480        assert_eq!(config.embedding_policies.len(), 2);
1481        assert!(config.api_key.is_none());
1482        assert!(config.api_base_url.is_some());
1483    }
1484
1485    #[test]
1486    fn test_nlq_config_construction() {
1487        let config = NLQConfig {
1488            enabled: false,
1489            provider: LLMProvider::AzureOpenAI,
1490            model: "gpt-4".to_string(),
1491            api_key: Some("azure-key".to_string()),
1492            api_base_url: Some("https://my-endpoint.openai.azure.com".to_string()),
1493            system_prompt: None,
1494        };
1495
1496        assert!(!config.enabled);
1497        assert_eq!(config.provider, LLMProvider::AzureOpenAI);
1498        assert!(config.api_key.is_some());
1499        assert!(config.api_base_url.is_some());
1500        assert!(config.system_prompt.is_none());
1501    }
1502
1503    #[test]
1504    fn test_agent_config_with_tools_and_policies() {
1505        let tool1 = ToolConfig {
1506            name: "cypher_query".to_string(),
1507            description: "Execute a Cypher query on the graph".to_string(),
1508            parameters: serde_json::json!({"query": {"type": "string"}, "readonly": {"type": "boolean"}}),
1509            enabled: true,
1510        };
1511        let tool2 = ToolConfig {
1512            name: "vector_search".to_string(),
1513            description: "Search vectors".to_string(),
1514            parameters: serde_json::json!({"query": {"type": "string"}, "k": {"type": "integer"}}),
1515            enabled: false,
1516        };
1517
1518        let mut policies = HashMap::new();
1519        policies.insert(
1520            "Person".to_string(),
1521            "Enrich person with external data".to_string(),
1522        );
1523        policies.insert("Company".to_string(), "Validate company info".to_string());
1524
1525        let config = AgentConfig {
1526            enabled: true,
1527            provider: LLMProvider::ClaudeCode,
1528            model: "claude-3-opus".to_string(),
1529            api_key: Some("sk-ant-key".to_string()),
1530            api_base_url: None,
1531            system_prompt: Some("You are a graph enrichment agent.".to_string()),
1532            tools: vec![tool1, tool2],
1533            policies,
1534        };
1535
1536        assert!(config.enabled);
1537        assert_eq!(config.tools.len(), 2);
1538        assert!(config.tools[0].enabled);
1539        assert!(!config.tools[1].enabled);
1540        assert_eq!(config.policies.len(), 2);
1541    }
1542
1543    #[test]
1544    fn test_create_tenant_with_custom_quotas() {
1545        let manager = TenantManager::new();
1546        let quotas = ResourceQuotas {
1547            max_nodes: Some(100),
1548            max_edges: Some(200),
1549            max_memory_bytes: Some(1024 * 1024),
1550            max_storage_bytes: None,
1551            max_connections: Some(5),
1552            max_query_time_ms: Some(10_000),
1553        };
1554
1555        manager
1556            .create_tenant("custom_t".to_string(), "Custom".to_string(), Some(quotas))
1557            .unwrap();
1558
1559        let tenant = manager.get_tenant("custom_t").unwrap();
1560        assert_eq!(tenant.quotas.max_nodes, Some(100));
1561        assert_eq!(tenant.quotas.max_edges, Some(200));
1562        assert!(tenant.quotas.max_storage_bytes.is_none());
1563    }
1564
1565    #[test]
1566    fn test_increment_decrement_all_resources() {
1567        let manager = TenantManager::new();
1568
1569        // Increment all resource types
1570        manager.increment_usage("default", "nodes", 10).unwrap();
1571        manager.increment_usage("default", "edges", 20).unwrap();
1572        manager.increment_usage("default", "memory", 1000).unwrap();
1573        manager.increment_usage("default", "storage", 2000).unwrap();
1574        manager
1575            .increment_usage("default", "connections", 5)
1576            .unwrap();
1577
1578        let usage = manager.get_usage("default").unwrap();
1579        assert_eq!(usage.node_count, 10);
1580        assert_eq!(usage.edge_count, 20);
1581        assert_eq!(usage.memory_bytes, 1000);
1582        assert_eq!(usage.storage_bytes, 2000);
1583        assert_eq!(usage.active_connections, 5);
1584
1585        // Decrement all
1586        manager.decrement_usage("default", "nodes", 3).unwrap();
1587        manager.decrement_usage("default", "edges", 5).unwrap();
1588        manager.decrement_usage("default", "memory", 200).unwrap();
1589        manager.decrement_usage("default", "storage", 500).unwrap();
1590        manager
1591            .decrement_usage("default", "connections", 2)
1592            .unwrap();
1593
1594        let usage = manager.get_usage("default").unwrap();
1595        assert_eq!(usage.node_count, 7);
1596        assert_eq!(usage.edge_count, 15);
1597        assert_eq!(usage.memory_bytes, 800);
1598        assert_eq!(usage.storage_bytes, 1500);
1599        assert_eq!(usage.active_connections, 3);
1600    }
1601
1602    #[test]
1603    fn test_tenant_serialization_with_configs() {
1604        let mut tenant = Tenant::new("full".to_string(), "Full Tenant".to_string());
1605        tenant.embed_config = Some(AutoEmbedConfig {
1606            provider: LLMProvider::OpenAI,
1607            embedding_model: "text-embedding-3-small".to_string(),
1608            api_key: None,
1609            api_base_url: None,
1610            chunk_size: 256,
1611            chunk_overlap: 32,
1612            vector_dimension: 1536,
1613            embedding_policies: HashMap::new(),
1614        });
1615        tenant.nlq_config = Some(NLQConfig {
1616            enabled: true,
1617            provider: LLMProvider::Mock,
1618            model: "mock".to_string(),
1619            api_key: None,
1620            api_base_url: None,
1621            system_prompt: None,
1622        });
1623
1624        let json = serde_json::to_string(&tenant).unwrap();
1625        let deserialized: Tenant = serde_json::from_str(&json).unwrap();
1626        assert!(deserialized.embed_config.is_some());
1627        assert!(deserialized.nlq_config.is_some());
1628        assert_eq!(deserialized.id, "full");
1629    }
1630
1631    #[test]
1632    fn test_resource_quotas_clone() {
1633        let quotas = ResourceQuotas {
1634            max_nodes: Some(42),
1635            max_edges: Some(84),
1636            max_memory_bytes: None,
1637            max_storage_bytes: None,
1638            max_connections: Some(10),
1639            max_query_time_ms: Some(5000),
1640        };
1641
1642        let cloned = quotas.clone();
1643        assert_eq!(cloned.max_nodes, Some(42));
1644        assert_eq!(cloned.max_edges, Some(84));
1645        assert_eq!(cloned.max_connections, Some(10));
1646    }
1647
1648    #[test]
1649    fn test_check_quota_storage_type_passes() {
1650        // Storage quota type is not explicitly enforced in check_quota match arms
1651        let manager = TenantManager::new();
1652        let result = manager.check_quota("default", "storage");
1653        assert!(result.is_ok());
1654    }
1655}
graphmind/persistence/tenant.rs

graphmind/persistence/
tenant.rs
