return old key on overwrite

src/lib.rs

@@ -23,16 +23,21 @@ impl Default for EthKeys {
 #[near_bindgen]
 impl EthKeys {
     // Sets the encrypted key for the sender's account.
-    pub fn set_key(&mut self, encrypted_key: String, overwrite: bool) {
+    pub fn set_key(&mut self, encrypted_key: String, overwrite: bool) -> Option<String> {
         // TODO - would be nice if there was some way to validate
         // that the encrypted key actualy contains expected data.
         let account_id = env::signer_account_id();
-        if !self.evm_key.contains_key(&account_id) || overwrite {
-            self.evm_key.insert(&account_id, &encrypted_key);
-            log_str(&format!("SetKey({account_id})"));
-        } else {
-            log_str(&format!("Key for {account_id} already set - use overwrite"));
-        }
+        let result = match self.evm_key.insert(&account_id, &encrypted_key) {
+            Some(old_encrypted_key) => {
+                if !overwrite {
+                    env::panic_str("Key already set!");
+                }
+                Some(old_encrypted_key)
+            }
+            None => None,
+        };
+        log_str(&format!("SetKey({account_id}, {encrypted_key})"));
+        result
     }
 
     pub fn get_key(&self, account_id: AccountId) -> Option<String> {
@@ -69,29 +74,34 @@ mod tests {
 
         let encrypted_key = "my_encrypted_key".to_string();
 
-        contract.set_key(encrypted_key.clone(), false);
-
-        let retrieved_key = contract.get_key(signer).unwrap();
+        let res = contract.set_key(encrypted_key.clone(), false);
+        assert_eq!(res, None);
 
-        assert_eq!(retrieved_key, encrypted_key);
+        let retrieved_key = contract.get_key(signer);
+        assert_eq!(Some(encrypted_key), retrieved_key);
     }
 
     #[test]
-    fn test_set_key_overwrite() {
+    #[should_panic]
+    fn test_set_key_overwrite_failure() {
         let signer = accounts(1);
         let context = get_context(signer.clone());
         let mut contract = EthKeys::default();
         testing_env!(context);
 
-        let encrypted_key = "my_encrypted_key".to_string();
-
-        contract.set_key(encrypted_key.clone(), false);
-
-        let new_key = "new_key".to_string();
-        contract.set_key(encrypted_key.clone(), false);
-        assert_eq!(contract.get_key(signer.clone()).unwrap(), encrypted_key);
+        contract.set_key("old_key".into(), false);
+        contract.set_key("new_key".into(), false);
+    }
 
-        contract.set_key(new_key.clone(), true);
-        assert_eq!(contract.get_key(signer).unwrap(), new_key);
+    #[test]
+    fn test_set_key_overwrite_passes() {
+        let signer = accounts(1);
+        let context = get_context(signer.clone());
+        let mut contract = EthKeys::default();
+        testing_env!(context);
+        let old_key = "old_key".to_string();
+        contract.set_key(old_key.clone(), false);
+        let reset_result = contract.set_key("new_key".into(), true);
+        assert_eq!(reset_result, Some(old_key));
     }
 }