Skip to main content

cc_switch/codex/
storage.rs

1use crate::codex::CodexConfiguration;
2use crate::config::ConfigStorage;
3use std::collections::BTreeMap;
4
5impl ConfigStorage {
6    /// Add a Codex configuration to storage
7    ///
8    /// Overwrites existing configuration with the same alias.
9    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    /// Get a Codex configuration by alias name
16    pub fn get_codex_configuration(&self, alias_name: &str) -> Option<&CodexConfiguration> {
17        self.codex_configurations.as_ref()?.get(alias_name)
18    }
19
20    /// Remove a Codex configuration by alias name
21    ///
22    /// Returns `true` if a configuration was found and removed, `false` otherwise.
23    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    /// Update a Codex configuration by old alias name
32    ///
33    /// If the alias changed, removes the old entry and inserts the new one.
34    /// Returns error if the old configuration doesn't exist.
35    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}