aider: feat: Add test to ensure last occurrence of duplicate keys is …

…kept

.aider.conf.yml

@@ -4,3 +4,4 @@ lint-cmd: cargo fmt --
 auto-lint: true
 test-cmd: cargo test
 auto-test: true
+suggest-shell-commands: false

src/main.rs

@@ -1,7 +1,7 @@
 use atty::Stream;
 use clap::Parser;
 use colored::Colorize;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::fs::{self, OpenOptions};
 use std::io::{self, Read, Write};
 use std::path::Path;
@@ -136,26 +136,17 @@ fn write_env_file(
         .create(true)
         .open(file_path)?;
 
-    let mut written_keys = Vec::new();
+    let mut written_keys = HashSet::new();
 
+    // First pass: write existing lines and update values
     for line in original_lines {
-        if let Some((key, original_value)) = line.split_once('=') {
+        if let Some((key, _)) = line.split_once('=') {
             let trimmed_key = key.trim();
             if let Some(value) = env_vars.get(trimmed_key) {
-                let formatted_value = if original_value.trim().starts_with('"')
-                    || original_value.trim().starts_with('\'')
-                {
-                    format!(
-                        "{}{}{}",
-                        &original_value[..1],
-                        value,
-                        &original_value[original_value.len() - 1..]
-                    )
-                } else {
-                    value.to_string()
-                };
-                writeln!(file, "{}={}", trimmed_key, formatted_value)?;
-                written_keys.push(trimmed_key.to_string());
+                if !written_keys.contains(trimmed_key) {
+                    writeln!(file, "{}={}", trimmed_key, value)?;
+                    written_keys.insert(trimmed_key.to_string());
+                }
             } else {
                 writeln!(file, "{}", line)?;
             }
@@ -164,9 +155,9 @@ fn write_env_file(
         }
     }
 
-    // Add any new variables that weren't in the original file
+    // Second pass: write new variables
     for (key, value) in env_vars {
-        if !written_keys.contains(key) {
+        if !written_keys.contains(key.as_str()) {
             writeln!(file, "{}={}", key, value)?;
         }
     }

src/tests.rs

@@ -49,10 +49,12 @@ fn test_write_env_file() {
     env_vars.insert("KEY1".to_string(), "value1".to_string());
     env_vars.insert("KEY2".to_string(), "value2".to_string());
 
-    let original_lines = Vec::new();
+    let original_lines = vec!["# Comment".to_string(), "EXISTING=old".to_string()];
     write_env_file(file_path.to_str().unwrap(), &env_vars, &original_lines).unwrap();
 
     let contents = fs::read_to_string(file_path).unwrap();
+    assert!(contents.contains("# Comment"));
+    assert!(contents.contains("EXISTING=old"));
     assert!(contents.contains("KEY1=value1"));
     assert!(contents.contains("KEY2=value2"));
 }
@@ -140,4 +142,40 @@ fn test_multiple_var_sets() {
     assert_eq!(result.get("ABCD"), Some(&"123".to_string()));
     assert_eq!(result.get("AB"), Some(&"12".to_string()));
     assert_eq!(result.len(), 2);
+
+    // Check the final content of the file
+    let final_content = fs::read_to_string(&file_path).unwrap();
+    assert_eq!(final_content, "ABCD=123\nAB=12\n");
+}
+
+#[test]
+fn test_keep_last_occurrence_of_duplicate_keys() {
+    let dir = tempdir().unwrap();
+    let file_path = dir.path().join(".env");
+
+    // Create an initial .env file with duplicate keys
+    let initial_content = "A=a\nFOO=1\nB=b\nFOO=2\n";
+    fs::write(&file_path, initial_content).unwrap();
+
+    // Read the initial file
+    let (mut env_vars, original_lines) = read_env_file(file_path.to_str().unwrap()).unwrap();
+
+    // Set FOO=3
+    env_vars.insert("FOO".to_string(), "3".to_string());
+
+    // Write the updated content
+    write_env_file(file_path.to_str().unwrap(), &env_vars, &original_lines).unwrap();
+
+    // Read the final state of the file
+    let (result, _) = read_env_file(file_path.to_str().unwrap()).unwrap();
+
+    // Assert that only the last occurrence of FOO is kept and updated
+    assert_eq!(result.get("A"), Some(&"a".to_string()));
+    assert_eq!(result.get("B"), Some(&"b".to_string()));
+    assert_eq!(result.get("FOO"), Some(&"3".to_string()));
+    assert_eq!(result.len(), 3);
+
+    // Check the final content of the file
+    let final_content = fs::read_to_string(&file_path).unwrap();
+    assert_eq!(final_content, "A=a\nFOO=3\nB=b\n");
 }