test: add comprehensive unit and integration tests for all modules

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
     rev: v2.4.2
     hooks:
       - id: codespell
-        entry: codespell -q 3 -f --skip=".git,.github,README.md,target,Cargo.lock" --ignore-words-list="astroid,braket,unstall,infinit,sems,te"
+        entry: codespell -q 3 -f --skip=".git,.github,README.md,target,Cargo.lock" --ignore-words-list="astroid,braket,unstall,infinit,sems,te,hel"
 
   - repo: https://github.com/jumanjihouse/pre-commit-hooks
     rev: 3.0.0

ares-cli/Cargo.toml

@@ -40,3 +40,4 @@ serde_yaml = "0.9"
 [dev-dependencies]
 tokio = { workspace = true }
 rstest = "0.26"
+ares-core = { path = "../ares-core", features = ["test-utils", "blue", "telemetry"] }

ares-cli/src/config.rs

@@ -305,3 +305,86 @@ fn replace_model_in_yaml(yaml: &str, role: &str, _old_model: &str, new_model: &s
 
     result
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn replace_model_basic() {
+        let yaml = "  orchestrator:\n    model: \"gpt-4\"\n    max_steps: 10\n";
+        let result = replace_model_in_yaml(yaml, "orchestrator", "gpt-4", "claude-3");
+        assert!(result.contains("model: \"claude-3\""));
+        assert!(!result.contains("gpt-4"));
+    }
+
+    #[test]
+    fn replace_model_preserves_other_roles() {
+        let yaml =
+            "  orchestrator:\n    model: \"gpt-4\"\n    max_steps: 10\n  recon:\n    model: \"gpt-4\"\n    max_steps: 5\n";
+        let result = replace_model_in_yaml(yaml, "orchestrator", "gpt-4", "claude-3");
+        // Only orchestrator should change
+        let lines: Vec<&str> = result.lines().collect();
+        let recon_idx = lines.iter().position(|l| l.contains("recon:")).unwrap();
+        let recon_model = lines[recon_idx + 1];
+        assert!(
+            recon_model.contains("gpt-4"),
+            "recon model should remain gpt-4"
+        );
+    }
+
+    #[test]
+    fn replace_model_role_not_found() {
+        let yaml = "  orchestrator:\n    model: \"gpt-4\"\n    max_steps: 10\n";
+        let result = replace_model_in_yaml(yaml, "nonexistent", "gpt-4", "claude-3");
+        assert_eq!(result, yaml);
+    }
+
+    #[test]
+    fn replace_model_preserves_indentation() {
+        let yaml = "  recon:\n    model: \"gpt-4\"\n    max_steps: 5\n";
+        let result = replace_model_in_yaml(yaml, "recon", "gpt-4", "claude-3");
+        assert!(result.contains("    model: \"claude-3\""));
+    }
+
+    #[test]
+    fn replace_model_no_trailing_newline() {
+        let yaml = "  recon:\n    model: \"gpt-4\"";
+        let result = replace_model_in_yaml(yaml, "recon", "gpt-4", "claude-3");
+        assert!(!result.ends_with('\n'));
+        assert!(result.contains("model: \"claude-3\""));
+    }
+
+    #[test]
+    fn replace_model_with_trailing_newline() {
+        let yaml = "  recon:\n    model: \"gpt-4\"\n";
+        let result = replace_model_in_yaml(yaml, "recon", "gpt-4", "claude-3");
+        assert!(result.ends_with('\n'));
+    }
+
+    #[test]
+    fn replace_model_preserves_surrounding_content() {
+        let yaml =
+            "# comment above\n  lateral:\n    model: \"old-model\"\n    max_steps: 20\n# comment below\n";
+        let result = replace_model_in_yaml(yaml, "lateral", "old-model", "new-model");
+        assert!(result.contains("# comment above"));
+        assert!(result.contains("# comment below"));
+        assert!(result.contains("    max_steps: 20"));
+    }
+
+    #[test]
+    fn replace_model_empty_yaml() {
+        let yaml = "";
+        let result = replace_model_in_yaml(yaml, "orchestrator", "gpt-4", "claude-3");
+        assert_eq!(result, "");
+    }
+
+    #[test]
+    fn replace_model_ignores_old_model_param() {
+        // The function uses _old_model (unused); it replaces whatever model: line
+        // is under the role, regardless of its current value.
+        let yaml = "  recon:\n    model: \"actual-model\"\n    max_steps: 5\n";
+        let result = replace_model_in_yaml(yaml, "recon", "wrong-model", "new-model");
+        assert!(result.contains("model: \"new-model\""));
+    }
+}

ares-cli/src/dedup/credentials.rs

@@ -100,3 +100,153 @@ pub(crate) fn dedup_credentials(creds: &[Credential]) -> Vec<Credential> {
     }
     result
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn make_cred(user: &str, pass: &str, domain: &str) -> Credential {
+        Credential {
+            id: uuid::Uuid::new_v4().to_string(),
+            username: user.to_string(),
+            password: pass.to_string(),
+            domain: domain.to_string(),
+            source: String::new(),
+            discovered_at: None,
+            is_admin: false,
+            parent_id: None,
+            attack_step: 0,
+        }
+    }
+
+    // ── strip_ansi ──────────────────────────────────────────────────
+
+    #[test]
+    fn strip_ansi_removes_color_codes() {
+        assert_eq!(strip_ansi("\x1b[31mred\x1b[0m"), "red");
+    }
+
+    #[test]
+    fn strip_ansi_passthrough_clean() {
+        assert_eq!(strip_ansi("clean text"), "clean text");
+    }
+
+    // ── sanitize_credentials ────────────────────────────────────────
+
+    #[test]
+    fn sanitize_strips_password_prefix() {
+        let mut creds = vec![make_cred("admin", "Password: Secret123", "contoso.local")];
+        sanitize_credentials(&mut creds);
+        assert_eq!(creds[0].password, "Secret123");
+    }
+
+    #[test]
+    fn sanitize_strips_trailing_paren() {
+        let mut creds = vec![make_cred("admin", "Secret123 (Pwn3d!)", "contoso.local")];
+        sanitize_credentials(&mut creds);
+        assert_eq!(creds[0].password, "Secret123");
+    }
+
+    #[test]
+    fn sanitize_removes_empty_password() {
+        let mut creds = vec![make_cred("admin", "", "contoso.local")];
+        sanitize_credentials(&mut creds);
+        assert!(creds.is_empty());
+    }
+
+    #[test]
+    fn sanitize_removes_password_literal() {
+        let mut creds = vec![make_cred("admin", "password", "contoso.local")];
+        sanitize_credentials(&mut creds);
+        assert!(creds.is_empty());
+    }
+
+    #[test]
+    fn sanitize_removes_discovered_marker() {
+        let mut creds = vec![make_cred("admin", "discovered", "contoso.local")];
+        sanitize_credentials(&mut creds);
+        assert!(creds.is_empty());
+    }
+
+    #[test]
+    fn sanitize_removes_hash_markers() {
+        let mut creds = vec![
+            make_cred("admin", "abc [NT]", "contoso.local"),
+            make_cred("admin", "def [SHA1]", "contoso.local"),
+        ];
+        sanitize_credentials(&mut creds);
+        assert!(creds.is_empty());
+    }
+
+    #[test]
+    fn sanitize_removes_slash_usernames() {
+        let mut creds = vec![make_cred("domain/admin", "pass", "contoso.local")];
+        sanitize_credentials(&mut creds);
+        assert!(creds.is_empty());
+    }
+
+    #[test]
+    fn sanitize_removes_evil_machine_accounts() {
+        let mut creds = vec![make_cred("evil$", "pass", "contoso.local")];
+        sanitize_credentials(&mut creds);
+        assert!(creds.is_empty());
+    }
+
+    #[test]
+    fn sanitize_extracts_domain_from_upn() {
+        let mut creds = vec![make_cred(
+            "sam.wilson@child.contoso.local",
+            "pass",
+            "old_domain",
+        )];
+        sanitize_credentials(&mut creds);
+        assert_eq!(creds[0].username, "sam.wilson");
+        assert_eq!(creds[0].domain, "child.contoso.local");
+    }
+
+    #[test]
+    fn sanitize_strips_trailing_dot_from_domain() {
+        let mut creds = vec![make_cred("admin", "pass", "contoso.local.")];
+        sanitize_credentials(&mut creds);
+        assert_eq!(creds[0].domain, "contoso.local");
+    }
+
+    // ── dedup_credentials ───────────────────────────────────────────
+
+    #[test]
+    fn dedup_removes_duplicates() {
+        let creds = vec![
+            make_cred("admin", "pass1", "contoso.local"),
+            make_cred("admin", "pass1", "contoso.local"),
+        ];
+        let result = dedup_credentials(&creds);
+        assert_eq!(result.len(), 1);
+    }
+
+    #[test]
+    fn dedup_keeps_different_passwords() {
+        let creds = vec![
+            make_cred("admin", "pass1", "contoso.local"),
+            make_cred("admin", "pass2", "contoso.local"),
+        ];
+        let result = dedup_credentials(&creds);
+        assert_eq!(result.len(), 2);
+    }
+
+    #[test]
+    fn dedup_skips_empty_passwords() {
+        let creds = vec![make_cred("admin", "", "contoso.local")];
+        let result = dedup_credentials(&creds);
+        assert!(result.is_empty());
+    }
+
+    #[test]
+    fn dedup_case_insensitive_key() {
+        let creds = vec![
+            make_cred("Admin", "pass1", "CONTOSO.LOCAL"),
+            make_cred("admin", "pass1", "contoso.local"),
+        ];
+        let result = dedup_credentials(&creds);
+        assert_eq!(result.len(), 1);
+    }
+}

ares-cli/src/dedup/labels.rs

@@ -100,3 +100,61 @@ pub(crate) fn normalize_source_label(source: &str) -> String {
         .collect::<Vec<_>>()
         .join(" ")
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn empty_source_returns_unknown() {
+        assert_eq!(normalize_source_label(""), "Unknown");
+    }
+
+    #[test]
+    fn exact_match_label() {
+        assert_eq!(normalize_source_label("recon"), "Reconnaissance");
+        assert_eq!(normalize_source_label("lateral"), "Lateral Movement");
+        assert_eq!(normalize_source_label("privesc"), "Privilege Escalation");
+        assert_eq!(normalize_source_label("crack"), "Password Cracking");
+    }
+
+    #[test]
+    fn case_insensitive_match() {
+        assert_eq!(normalize_source_label("RECON"), "Reconnaissance");
+        assert_eq!(normalize_source_label("Exploit"), "Exploitation");
+    }
+
+    #[test]
+    fn dedup_colon_prefix() {
+        assert_eq!(normalize_source_label("recon:recon"), "Reconnaissance");
+    }
+
+    #[test]
+    fn task_input_pattern_extracts_type() {
+        assert_eq!(
+            normalize_source_label("task input (recon_abc12345)"),
+            "Reconnaissance"
+        );
+    }
+
+    #[test]
+    fn task_suffix_strips_id() {
+        assert_eq!(
+            normalize_source_label("recon_abc12345678"),
+            "Reconnaissance"
+        );
+    }
+
+    #[test]
+    fn fallback_title_cases() {
+        let result = normalize_source_label("some_custom_source");
+        assert_eq!(result, "Some Custom Source");
+    }
+
+    #[test]
+    fn tool_based_sources() {
+        assert_eq!(normalize_source_label("secretsdump"), "Secretsdump");
+        assert_eq!(normalize_source_label("kerberoast"), "Kerberoasting");
+        assert_eq!(normalize_source_label("bloodhound"), "BloodHound");
+    }
+}