From b23a707b11b5d1aaf12ad61add83699a5324e73e Mon Sep 17 00:00:00 2001
From: Andrew Lewis <nerf@judo.za.org>
Date: Tue, 30 Jul 2024 16:16:28 +0200
Subject: [PATCH] [Test] Add more tests for milter_headers

---
 .../functional/cases/550_milter_headers.robot | 21 +++++++++++++++++-
 test/functional/configs/milter_headers.conf   | 11 +++++++++-
 test/functional/lib/rspamd.robot              | 22 +++++++++++++++++++
 3 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/test/functional/cases/550_milter_headers.robot b/test/functional/cases/550_milter_headers.robot
index 737d0b42d..80471b83c 100644
--- a/test/functional/cases/550_milter_headers.robot
+++ b/test/functional/cases/550_milter_headers.robot
@@ -11,10 +11,29 @@ ${MESSAGE}               ${RSPAMD_TESTDIR}/messages/zip.eml
 ${RSPAMD_SCOPE}          Suite
 ${RSPAMD_URL_TLD}        ${RSPAMD_TESTDIR}/../lua/unit/test_tld.dat
 ${SETTINGS_NOSYMBOLS}    {symbols_enabled = []}
-${SETTINGS_TEST}         {symbols_enabled = [SIMPLE_TEST]}
+${SETTINGS_TEST}         {SIMPLE_TEST = 2.0, symbols_enabled = [SIMPLE_TEST]}
 
 *** Test Cases ***
 CHECK HEADERS WITH TEST SYMBOL
   Scan File  ${MESSAGE}  Settings=${SETTINGS_TEST}
+  # Check X-Virus header
   Expect Removed Header  X-Virus
   Expect Added Header  X-Virus  Fires always
+  # Check My-Spamd-Bar header
+  Expect Added Header  My-Spamd-Bar  ++
+  Do Not Expect Removed Header  My-Spamd-Bar
+  # Check X-Spam-Level header
+  Expect Added Header  X-Spam-Level  **
+  Expect Removed Header  X-Spam-Level
+
+CHECK HEADERS WITHOUT TEST SYMBOL
+  Scan File  ${MESSAGE}  Settings=${SETTINGS_NOSYMBOLS}
+  # Check X-Virus header
+  Expect Removed Header  X-Virus
+  Do Not Expect Added Header  X-Virus
+  # Check My-Spamd-Bar header
+  Expect Added Header  My-Spamd-Bar  /
+  Do Not Expect Removed Header  My-Spamd-Bar
+  # Check X-Spam-Level header
+  Do Not Expect Added Header  X-Spam-Level
+  Expect Removed Header  X-Spam-Level
diff --git a/test/functional/configs/milter_headers.conf b/test/functional/configs/milter_headers.conf
index 0787e9463..947bc28dd 100644
--- a/test/functional/configs/milter_headers.conf
+++ b/test/functional/configs/milter_headers.conf
@@ -4,9 +4,18 @@ lua = "{= env.TESTDIR =}/lua/simple.lua"
 
 milter_headers {
 
-  use = ["x-virus"];
+  use = ["remove-headers", "x-spam-level", "x-spamd-bar", "x-virus"];
 
   routines {
+    remove-headers {
+      headers {
+        "X-Spam-Level" = 0,
+      }
+    }
+    x-spamd-bar {
+      header = "My-Spamd-Bar";
+      remove = null;
+    }
     x-virus {
       symbols = ["SIMPLE_TEST"];
     }
diff --git a/test/functional/lib/rspamd.robot b/test/functional/lib/rspamd.robot
index 917b01298..c45a19908 100644
--- a/test/functional/lib/rspamd.robot
+++ b/test/functional/lib/rspamd.robot
@@ -60,6 +60,28 @@ Check Rspamc Match String
     Should Not Contain  ${subject}  ${str}
   END
 
+Do Not Expect Added Header
+  [Arguments]  ${header_name}
+  IF  'milter' not in ${SCAN_RESULT}
+    RETURN
+  END
+  IF  'add_headers' not in ${SCAN_RESULT}[milter]
+    RETURN
+  END
+  Dictionary Should Not Contain Key  ${SCAN_RESULT}[milter][add_headers]  ${header_name}
+  ...  msg=${header_name} was added
+
+Do Not Expect Removed Header
+  [Arguments]  ${header_name}
+  IF  'milter' not in ${SCAN_RESULT}
+    RETURN
+  END
+  IF  'remove_headers' not in ${SCAN_RESULT}[milter]
+    RETURN
+  END
+  Dictionary Should Not Contain Key  ${SCAN_RESULT}[milter][remove_headers]  ${header_name}
+  ...  msg=${header_name} was removed
+
 Do Not Expect Symbol
   [Arguments]  ${symbol}
   Dictionary Should Not Contain Key  ${SCAN_RESULT}[symbols]  ${symbol}
-- 
2.39.5