litellm_rs/core/virtual_keys/
requests.rs

1//! Virtual key request types
2
3use super::types::{Permission, RateLimits};
4use chrono::{DateTime, Utc};
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7
8/// Virtual key creation request
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct CreateKeyRequest {
11    /// Key alias (optional)
12    pub key_alias: Option<String>,
13    /// User ID
14    pub user_id: String,
15    /// Team ID (optional)
16    pub team_id: Option<String>,
17    /// Models to allow
18    pub models: Vec<String>,
19    /// Maximum budget
20    pub max_budget: Option<f64>,
21    /// Budget duration
22    pub budget_duration: Option<String>,
23    /// Rate limits
24    pub rate_limits: Option<RateLimits>,
25    /// Permissions
26    pub permissions: Vec<Permission>,
27    /// Metadata
28    pub metadata: HashMap<String, String>,
29    /// Expiration time
30    pub expires_at: Option<DateTime<Utc>>,
31    /// Tags
32    pub tags: Vec<String>,
33}
34
35/// Virtual key update request
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct UpdateKeyRequest {
38    /// Key alias
39    pub key_alias: Option<String>,
40    /// Models to allow
41    pub models: Option<Vec<String>>,
42    /// Maximum budget
43    pub max_budget: Option<f64>,
44    /// Budget duration
45    pub budget_duration: Option<String>,
46    /// Rate limits
47    pub rate_limits: Option<RateLimits>,
48    /// Permissions
49    pub permissions: Option<Vec<Permission>>,
50    /// Metadata
51    pub metadata: Option<HashMap<String, String>>,
52    /// Expiration time
53    pub expires_at: Option<DateTime<Utc>>,
54    /// Whether key is active
55    pub is_active: Option<bool>,
56    /// Tags
57    pub tags: Option<Vec<String>>,
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63    use chrono::{Duration, Utc};
64
65    // ==================== Helper Functions ====================
66
67    fn create_test_rate_limits() -> RateLimits {
68        RateLimits {
69            rpm: Some(60),
70            rph: Some(3600),
71            rpd: Some(86400),
72            tpm: Some(100000),
73            tph: Some(6000000),
74            tpd: Some(144000000),
75            max_parallel_requests: Some(10),
76        }
77    }
78
79    // ==================== CreateKeyRequest Tests ====================
80
81    #[test]
82    fn test_create_key_request_minimal() {
83        let request = CreateKeyRequest {
84            key_alias: None,
85            user_id: "user-001".to_string(),
86            team_id: None,
87            models: vec![],
88            max_budget: None,
89            budget_duration: None,
90            rate_limits: None,
91            permissions: vec![],
92            metadata: HashMap::new(),
93            expires_at: None,
94            tags: vec![],
95        };
96
97        assert_eq!(request.user_id, "user-001");
98        assert!(request.key_alias.is_none());
99        assert!(request.models.is_empty());
100        assert!(request.permissions.is_empty());
101    }
102
103    #[test]
104    fn test_create_key_request_full() {
105        let expires = Utc::now() + Duration::days(365);
106        let mut metadata = HashMap::new();
107        metadata.insert("env".to_string(), "production".to_string());
108
109        let request = CreateKeyRequest {
110            key_alias: Some("my-api-key".to_string()),
111            user_id: "user-001".to_string(),
112            team_id: Some("team-001".to_string()),
113            models: vec!["gpt-4".to_string(), "gpt-3.5-turbo".to_string()],
114            max_budget: Some(100.0),
115            budget_duration: Some("1m".to_string()),
116            rate_limits: Some(create_test_rate_limits()),
117            permissions: vec![Permission::ChatCompletion, Permission::Embedding],
118            metadata,
119            expires_at: Some(expires),
120            tags: vec!["production".to_string(), "api".to_string()],
121        };
122
123        assert_eq!(request.key_alias, Some("my-api-key".to_string()));
124        assert_eq!(request.user_id, "user-001");
125        assert_eq!(request.team_id, Some("team-001".to_string()));
126        assert_eq!(request.models.len(), 2);
127        assert_eq!(request.max_budget, Some(100.0));
128        assert_eq!(request.permissions.len(), 2);
129        assert_eq!(request.tags.len(), 2);
130    }
131
132    #[test]
133    fn test_create_key_request_with_models() {
134        let request = CreateKeyRequest {
135            key_alias: None,
136            user_id: "user-001".to_string(),
137            team_id: None,
138            models: vec![
139                "gpt-4".to_string(),
140                "gpt-3.5-turbo".to_string(),
141                "claude-3-opus".to_string(),
142            ],
143            max_budget: None,
144            budget_duration: None,
145            rate_limits: None,
146            permissions: vec![],
147            metadata: HashMap::new(),
148            expires_at: None,
149            tags: vec![],
150        };
151
152        assert_eq!(request.models.len(), 3);
153        assert!(request.models.contains(&"gpt-4".to_string()));
154        assert!(request.models.contains(&"claude-3-opus".to_string()));
155    }
156
157    #[test]
158    fn test_create_key_request_with_budget() {
159        let request = CreateKeyRequest {
160            key_alias: None,
161            user_id: "user-001".to_string(),
162            team_id: None,
163            models: vec![],
164            max_budget: Some(500.0),
165            budget_duration: Some("1w".to_string()),
166            rate_limits: None,
167            permissions: vec![],
168            metadata: HashMap::new(),
169            expires_at: None,
170            tags: vec![],
171        };
172
173        assert_eq!(request.max_budget, Some(500.0));
174        assert_eq!(request.budget_duration, Some("1w".to_string()));
175    }
176
177    #[test]
178    fn test_create_key_request_with_rate_limits() {
179        let request = CreateKeyRequest {
180            key_alias: None,
181            user_id: "user-001".to_string(),
182            team_id: None,
183            models: vec![],
184            max_budget: None,
185            budget_duration: None,
186            rate_limits: Some(create_test_rate_limits()),
187            permissions: vec![],
188            metadata: HashMap::new(),
189            expires_at: None,
190            tags: vec![],
191        };
192
193        assert!(request.rate_limits.is_some());
194        let limits = request.rate_limits.unwrap();
195        assert_eq!(limits.rpm, Some(60));
196        assert_eq!(limits.tpm, Some(100000));
197    }
198
199    #[test]
200    fn test_create_key_request_with_permissions() {
201        let request = CreateKeyRequest {
202            key_alias: None,
203            user_id: "user-001".to_string(),
204            team_id: None,
205            models: vec![],
206            max_budget: None,
207            budget_duration: None,
208            rate_limits: None,
209            permissions: vec![
210                Permission::ChatCompletion,
211                Permission::TextCompletion,
212                Permission::Embedding,
213                Permission::Admin,
214            ],
215            metadata: HashMap::new(),
216            expires_at: None,
217            tags: vec![],
218        };
219
220        assert_eq!(request.permissions.len(), 4);
221        assert!(request.permissions.contains(&Permission::Admin));
222    }
223
224    #[test]
225    fn test_create_key_request_with_metadata() {
226        let mut metadata = HashMap::new();
227        metadata.insert("environment".to_string(), "production".to_string());
228        metadata.insert("version".to_string(), "v1".to_string());
229        metadata.insert("owner".to_string(), "team-a".to_string());
230
231        let request = CreateKeyRequest {
232            key_alias: None,
233            user_id: "user-001".to_string(),
234            team_id: None,
235            models: vec![],
236            max_budget: None,
237            budget_duration: None,
238            rate_limits: None,
239            permissions: vec![],
240            metadata,
241            expires_at: None,
242            tags: vec![],
243        };
244
245        assert_eq!(request.metadata.len(), 3);
246        assert_eq!(
247            request.metadata.get("environment"),
248            Some(&"production".to_string())
249        );
250    }
251
252    #[test]
253    fn test_create_key_request_with_expiration() {
254        let expires = Utc::now() + Duration::days(30);
255
256        let request = CreateKeyRequest {
257            key_alias: None,
258            user_id: "user-001".to_string(),
259            team_id: None,
260            models: vec![],
261            max_budget: None,
262            budget_duration: None,
263            rate_limits: None,
264            permissions: vec![],
265            metadata: HashMap::new(),
266            expires_at: Some(expires),
267            tags: vec![],
268        };
269
270        assert!(request.expires_at.is_some());
271        assert!(request.expires_at.unwrap() > Utc::now());
272    }
273
274    #[test]
275    fn test_create_key_request_clone() {
276        let request = CreateKeyRequest {
277            key_alias: Some("test-key".to_string()),
278            user_id: "user-001".to_string(),
279            team_id: Some("team-001".to_string()),
280            models: vec!["gpt-4".to_string()],
281            max_budget: Some(100.0),
282            budget_duration: Some("1m".to_string()),
283            rate_limits: None,
284            permissions: vec![Permission::ChatCompletion],
285            metadata: HashMap::new(),
286            expires_at: None,
287            tags: vec!["test".to_string()],
288        };
289
290        let cloned = request.clone();
291        assert_eq!(cloned.key_alias, request.key_alias);
292        assert_eq!(cloned.user_id, request.user_id);
293        assert_eq!(cloned.models, request.models);
294    }
295
296    #[test]
297    fn test_create_key_request_debug() {
298        let request = CreateKeyRequest {
299            key_alias: Some("debug-key".to_string()),
300            user_id: "user-001".to_string(),
301            team_id: None,
302            models: vec![],
303            max_budget: None,
304            budget_duration: None,
305            rate_limits: None,
306            permissions: vec![],
307            metadata: HashMap::new(),
308            expires_at: None,
309            tags: vec![],
310        };
311
312        let debug_str = format!("{:?}", request);
313        assert!(debug_str.contains("CreateKeyRequest"));
314        assert!(debug_str.contains("debug-key"));
315    }
316
317    #[test]
318    fn test_create_key_request_serialization() {
319        let request = CreateKeyRequest {
320            key_alias: Some("ser-key".to_string()),
321            user_id: "user-001".to_string(),
322            team_id: None,
323            models: vec!["gpt-4".to_string()],
324            max_budget: Some(100.0),
325            budget_duration: None,
326            rate_limits: None,
327            permissions: vec![Permission::ChatCompletion],
328            metadata: HashMap::new(),
329            expires_at: None,
330            tags: vec![],
331        };
332
333        let json = serde_json::to_string(&request).unwrap();
334        assert!(json.contains("ser-key"));
335        assert!(json.contains("user-001"));
336        assert!(json.contains("gpt-4"));
337
338        let parsed: CreateKeyRequest = serde_json::from_str(&json).unwrap();
339        assert_eq!(parsed.key_alias, request.key_alias);
340        assert_eq!(parsed.user_id, request.user_id);
341    }
342
343    // ==================== UpdateKeyRequest Tests ====================
344
345    #[test]
346    fn test_update_key_request_minimal() {
347        let request = UpdateKeyRequest {
348            key_alias: None,
349            models: None,
350            max_budget: None,
351            budget_duration: None,
352            rate_limits: None,
353            permissions: None,
354            metadata: None,
355            expires_at: None,
356            is_active: None,
357            tags: None,
358        };
359
360        assert!(request.key_alias.is_none());
361        assert!(request.models.is_none());
362        assert!(request.is_active.is_none());
363    }
364
365    #[test]
366    fn test_update_key_request_full() {
367        let expires = Utc::now() + Duration::days(90);
368        let mut metadata = HashMap::new();
369        metadata.insert("updated".to_string(), "true".to_string());
370
371        let request = UpdateKeyRequest {
372            key_alias: Some("updated-key".to_string()),
373            models: Some(vec!["gpt-4".to_string()]),
374            max_budget: Some(200.0),
375            budget_duration: Some("1m".to_string()),
376            rate_limits: Some(create_test_rate_limits()),
377            permissions: Some(vec![Permission::ChatCompletion]),
378            metadata: Some(metadata),
379            expires_at: Some(expires),
380            is_active: Some(true),
381            tags: Some(vec!["updated".to_string()]),
382        };
383
384        assert_eq!(request.key_alias, Some("updated-key".to_string()));
385        assert_eq!(request.max_budget, Some(200.0));
386        assert_eq!(request.is_active, Some(true));
387    }
388
389    #[test]
390    fn test_update_key_request_partial_update() {
391        let request = UpdateKeyRequest {
392            key_alias: Some("new-alias".to_string()),
393            models: None,
394            max_budget: Some(300.0),
395            budget_duration: None,
396            rate_limits: None,
397            permissions: None,
398            metadata: None,
399            expires_at: None,
400            is_active: None,
401            tags: None,
402        };
403
404        assert_eq!(request.key_alias, Some("new-alias".to_string()));
405        assert!(request.models.is_none());
406        assert_eq!(request.max_budget, Some(300.0));
407    }
408
409    #[test]
410    fn test_update_key_request_deactivate() {
411        let request = UpdateKeyRequest {
412            key_alias: None,
413            models: None,
414            max_budget: None,
415            budget_duration: None,
416            rate_limits: None,
417            permissions: None,
418            metadata: None,
419            expires_at: None,
420            is_active: Some(false),
421            tags: None,
422        };
423
424        assert_eq!(request.is_active, Some(false));
425    }
426
427    #[test]
428    fn test_update_key_request_update_models() {
429        let request = UpdateKeyRequest {
430            key_alias: None,
431            models: Some(vec![
432                "gpt-4".to_string(),
433                "gpt-4-turbo".to_string(),
434                "claude-3-opus".to_string(),
435            ]),
436            max_budget: None,
437            budget_duration: None,
438            rate_limits: None,
439            permissions: None,
440            metadata: None,
441            expires_at: None,
442            is_active: None,
443            tags: None,
444        };
445
446        assert!(request.models.is_some());
447        let models = request.models.unwrap();
448        assert_eq!(models.len(), 3);
449    }
450
451    #[test]
452    fn test_update_key_request_update_rate_limits() {
453        let request = UpdateKeyRequest {
454            key_alias: None,
455            models: None,
456            max_budget: None,
457            budget_duration: None,
458            rate_limits: Some(RateLimits {
459                rpm: Some(120),
460                rph: None,
461                rpd: None,
462                tpm: Some(200000),
463                tph: None,
464                tpd: None,
465                max_parallel_requests: Some(20),
466            }),
467            permissions: None,
468            metadata: None,
469            expires_at: None,
470            is_active: None,
471            tags: None,
472        };
473
474        assert!(request.rate_limits.is_some());
475        let limits = request.rate_limits.unwrap();
476        assert_eq!(limits.rpm, Some(120));
477        assert_eq!(limits.tpm, Some(200000));
478    }
479
480    #[test]
481    fn test_update_key_request_update_permissions() {
482        let request = UpdateKeyRequest {
483            key_alias: None,
484            models: None,
485            max_budget: None,
486            budget_duration: None,
487            rate_limits: None,
488            permissions: Some(vec![
489                Permission::ChatCompletion,
490                Permission::Admin,
491                Permission::KeyManagement,
492            ]),
493            metadata: None,
494            expires_at: None,
495            is_active: None,
496            tags: None,
497        };
498
499        assert!(request.permissions.is_some());
500        let perms = request.permissions.unwrap();
501        assert_eq!(perms.len(), 3);
502        assert!(perms.contains(&Permission::Admin));
503    }
504
505    #[test]
506    fn test_update_key_request_update_metadata() {
507        let mut metadata = HashMap::new();
508        metadata.insert("status".to_string(), "upgraded".to_string());
509        metadata.insert("tier".to_string(), "premium".to_string());
510
511        let request = UpdateKeyRequest {
512            key_alias: None,
513            models: None,
514            max_budget: None,
515            budget_duration: None,
516            rate_limits: None,
517            permissions: None,
518            metadata: Some(metadata),
519            expires_at: None,
520            is_active: None,
521            tags: None,
522        };
523
524        assert!(request.metadata.is_some());
525        let meta = request.metadata.unwrap();
526        assert_eq!(meta.len(), 2);
527        assert_eq!(meta.get("tier"), Some(&"premium".to_string()));
528    }
529
530    #[test]
531    fn test_update_key_request_extend_expiration() {
532        let new_expires = Utc::now() + Duration::days(365);
533
534        let request = UpdateKeyRequest {
535            key_alias: None,
536            models: None,
537            max_budget: None,
538            budget_duration: None,
539            rate_limits: None,
540            permissions: None,
541            metadata: None,
542            expires_at: Some(new_expires),
543            is_active: None,
544            tags: None,
545        };
546
547        assert!(request.expires_at.is_some());
548        assert!(request.expires_at.unwrap() > Utc::now());
549    }
550
551    #[test]
552    fn test_update_key_request_clone() {
553        let request = UpdateKeyRequest {
554            key_alias: Some("clone-test".to_string()),
555            models: Some(vec!["gpt-4".to_string()]),
556            max_budget: Some(100.0),
557            budget_duration: None,
558            rate_limits: None,
559            permissions: None,
560            metadata: None,
561            expires_at: None,
562            is_active: Some(true),
563            tags: None,
564        };
565
566        let cloned = request.clone();
567        assert_eq!(cloned.key_alias, request.key_alias);
568        assert_eq!(cloned.max_budget, request.max_budget);
569        assert_eq!(cloned.is_active, request.is_active);
570    }
571
572    #[test]
573    fn test_update_key_request_debug() {
574        let request = UpdateKeyRequest {
575            key_alias: Some("debug-test".to_string()),
576            models: None,
577            max_budget: None,
578            budget_duration: None,
579            rate_limits: None,
580            permissions: None,
581            metadata: None,
582            expires_at: None,
583            is_active: None,
584            tags: None,
585        };
586
587        let debug_str = format!("{:?}", request);
588        assert!(debug_str.contains("UpdateKeyRequest"));
589        assert!(debug_str.contains("debug-test"));
590    }
591
592    #[test]
593    fn test_update_key_request_serialization() {
594        let request = UpdateKeyRequest {
595            key_alias: Some("ser-test".to_string()),
596            models: Some(vec!["gpt-4".to_string()]),
597            max_budget: Some(150.0),
598            budget_duration: None,
599            rate_limits: None,
600            permissions: None,
601            metadata: None,
602            expires_at: None,
603            is_active: Some(true),
604            tags: None,
605        };
606
607        let json = serde_json::to_string(&request).unwrap();
608        assert!(json.contains("ser-test"));
609        assert!(json.contains("150"));
610
611        let parsed: UpdateKeyRequest = serde_json::from_str(&json).unwrap();
612        assert_eq!(parsed.key_alias, request.key_alias);
613        assert_eq!(parsed.max_budget, request.max_budget);
614    }
615
616    // ==================== Integration Tests ====================
617
618    #[test]
619    fn test_key_lifecycle_create_and_update() {
620        // Step 1: Create key request
621        let create_request = CreateKeyRequest {
622            key_alias: Some("lifecycle-key".to_string()),
623            user_id: "user-001".to_string(),
624            team_id: None,
625            models: vec!["gpt-3.5-turbo".to_string()],
626            max_budget: Some(50.0),
627            budget_duration: None,
628            rate_limits: None,
629            permissions: vec![Permission::ChatCompletion],
630            metadata: HashMap::new(),
631            expires_at: None,
632            tags: vec!["test".to_string()],
633        };
634
635        assert_eq!(create_request.max_budget, Some(50.0));
636        assert_eq!(create_request.models.len(), 1);
637
638        // Step 2: Upgrade key
639        let update_request = UpdateKeyRequest {
640            key_alias: None,
641            models: Some(vec!["gpt-3.5-turbo".to_string(), "gpt-4".to_string()]),
642            max_budget: Some(200.0),
643            budget_duration: None,
644            rate_limits: Some(create_test_rate_limits()),
645            permissions: Some(vec![Permission::ChatCompletion, Permission::Embedding]),
646            metadata: None,
647            expires_at: None,
648            is_active: None,
649            tags: None,
650        };
651
652        assert_eq!(update_request.max_budget, Some(200.0));
653        assert_eq!(update_request.models.as_ref().unwrap().len(), 2);
654    }
655
656    #[test]
657    fn test_budget_duration_formats() {
658        let durations = vec!["1d", "1w", "1m", "3m", "1y"];
659
660        for duration in durations {
661            let request = CreateKeyRequest {
662                key_alias: None,
663                user_id: "user-001".to_string(),
664                team_id: None,
665                models: vec![],
666                max_budget: Some(100.0),
667                budget_duration: Some(duration.to_string()),
668                rate_limits: None,
669                permissions: vec![],
670                metadata: HashMap::new(),
671                expires_at: None,
672                tags: vec![],
673            };
674
675            assert_eq!(request.budget_duration, Some(duration.to_string()));
676        }
677    }
678
679    #[test]
680    fn test_permission_combinations() {
681        let basic_perms = [Permission::ChatCompletion];
682        let standard_perms = [
683            Permission::ChatCompletion,
684            Permission::TextCompletion,
685            Permission::Embedding,
686        ];
687        let admin_perms = [
688            Permission::ChatCompletion,
689            Permission::TextCompletion,
690            Permission::Embedding,
691            Permission::Admin,
692            Permission::KeyManagement,
693        ];
694
695        assert_eq!(basic_perms.len(), 1);
696        assert_eq!(standard_perms.len(), 3);
697        assert_eq!(admin_perms.len(), 5);
698        assert!(admin_perms.contains(&Permission::Admin));
699    }
700
701    #[test]
702    fn test_empty_update_request() {
703        let request = UpdateKeyRequest {
704            key_alias: None,
705            models: None,
706            max_budget: None,
707            budget_duration: None,
708            rate_limits: None,
709            permissions: None,
710            metadata: None,
711            expires_at: None,
712            is_active: None,
713            tags: None,
714        };
715
716        // All fields are None - valid empty update
717        assert!(request.key_alias.is_none());
718        assert!(request.models.is_none());
719        assert!(request.max_budget.is_none());
720        assert!(request.budget_duration.is_none());
721        assert!(request.rate_limits.is_none());
722        assert!(request.permissions.is_none());
723        assert!(request.metadata.is_none());
724        assert!(request.expires_at.is_none());
725        assert!(request.is_active.is_none());
726        assert!(request.tags.is_none());
727    }
728
729    #[test]
730    fn test_update_tags() {
731        let request = UpdateKeyRequest {
732            key_alias: None,
733            models: None,
734            max_budget: None,
735            budget_duration: None,
736            rate_limits: None,
737            permissions: None,
738            metadata: None,
739            expires_at: None,
740            is_active: None,
741            tags: Some(vec!["production".to_string(), "high-priority".to_string()]),
742        };
743
744        assert!(request.tags.is_some());
745        let tags = request.tags.unwrap();
746        assert_eq!(tags.len(), 2);
747        assert!(tags.contains(&"production".to_string()));
748    }
749}
litellm_rs/core/virtual_keys/requests.rs

litellm_rs/core/virtual_keys/
requests.rs
