cc_switch/codex/
storage.rs1use crate::codex::CodexConfiguration;
2use crate::config::ConfigStorage;
3use std::collections::BTreeMap;
4
5impl ConfigStorage {
6 pub fn add_codex_configuration(&mut self, config: CodexConfiguration) {
10 self.codex_configurations
11 .get_or_insert_with(BTreeMap::new)
12 .insert(config.alias_name.clone(), config);
13 }
14
15 pub fn get_codex_configuration(&self, alias_name: &str) -> Option<&CodexConfiguration> {
17 self.codex_configurations.as_ref()?.get(alias_name)
18 }
19
20 pub fn remove_codex_configuration(&mut self, alias_name: &str) -> bool {
24 if let Some(ref mut map) = self.codex_configurations {
25 map.remove(alias_name).is_some()
26 } else {
27 false
28 }
29 }
30
31 pub fn update_codex_configuration(
36 &mut self,
37 old_alias: &str,
38 new_config: CodexConfiguration,
39 ) -> anyhow::Result<()> {
40 let map = self
41 .codex_configurations
42 .as_mut()
43 .ok_or_else(|| anyhow::anyhow!("No Codex configurations in storage"))?;
44
45 if !map.contains_key(old_alias) {
46 return Err(anyhow::anyhow!(
47 "Codex configuration '{}' not found",
48 old_alias
49 ));
50 }
51
52 if old_alias != new_config.alias_name {
53 map.remove(old_alias);
54 }
55
56 map.insert(new_config.alias_name.clone(), new_config);
57 Ok(())
58 }
59}
60
61#[cfg(test)]
62mod tests {
63 use super::*;
64
65 fn make_apikey_config(alias: &str) -> CodexConfiguration {
66 CodexConfiguration {
67 alias_name: alias.to_string(),
68 auth_mode: "apikey".to_string(),
69 openai_api_key: Some(format!("sk-{alias}")),
70 id_token: None,
71 access_token: None,
72 refresh_token: None,
73 account_id: None,
74 last_refresh: None,
75 }
76 }
77
78 #[test]
79 fn test_add_codex_configuration() {
80 let mut storage = ConfigStorage::default();
81 assert!(storage.codex_configurations.is_none());
82
83 let config = make_apikey_config("test");
84 storage.add_codex_configuration(config.clone());
85
86 assert!(storage.codex_configurations.is_some());
87 assert_eq!(storage.codex_configurations.as_ref().unwrap().len(), 1);
88 assert_eq!(
89 storage.codex_configurations.as_ref().unwrap().get("test"),
90 Some(&config)
91 );
92 }
93
94 #[test]
95 fn test_add_overwrites_existing() {
96 let mut storage = ConfigStorage::default();
97 storage.add_codex_configuration(make_apikey_config("test"));
98
99 let mut updated = make_apikey_config("test");
100 updated.openai_api_key = Some("sk-updated".to_string());
101 storage.add_codex_configuration(updated.clone());
102
103 assert_eq!(storage.codex_configurations.as_ref().unwrap().len(), 1);
104 assert_eq!(
105 storage
106 .get_codex_configuration("test")
107 .unwrap()
108 .openai_api_key,
109 Some("sk-updated".to_string())
110 );
111 }
112
113 #[test]
114 fn test_get_codex_configuration() {
115 let mut storage = ConfigStorage::default();
116 let config = make_apikey_config("foo");
117 storage.add_codex_configuration(config.clone());
118
119 assert_eq!(storage.get_codex_configuration("foo"), Some(&config));
120 assert_eq!(storage.get_codex_configuration("missing"), None);
121 }
122
123 #[test]
124 fn test_get_codex_configuration_empty_storage() {
125 let storage = ConfigStorage::default();
126 assert_eq!(storage.get_codex_configuration("any"), None);
127 }
128
129 #[test]
130 fn test_remove_codex_configuration() {
131 let mut storage = ConfigStorage::default();
132 storage.add_codex_configuration(make_apikey_config("test"));
133
134 assert!(storage.remove_codex_configuration("test"));
135 assert_eq!(storage.get_codex_configuration("test"), None);
136 }
137
138 #[test]
139 fn test_remove_nonexistent_codex_configuration() {
140 let mut storage = ConfigStorage::default();
141 storage.add_codex_configuration(make_apikey_config("test"));
142
143 assert!(!storage.remove_codex_configuration("missing"));
144 assert_eq!(storage.codex_configurations.as_ref().unwrap().len(), 1);
145 }
146
147 #[test]
148 fn test_remove_from_empty_storage() {
149 let mut storage = ConfigStorage::default();
150 assert!(!storage.remove_codex_configuration("any"));
151 }
152}