Merge pull request #5083 from fatalbanana/test_milter_headers

[Test] Add tests for milter_headers

3 files changed, 106 insertions, 0 deletions

diff --git a/test/functional/cases/550_milter_headers.robot b/test/functional/cases/550_milter_headers.robot
new file mode 100644
index 000000000..80471b83c
--- /dev/null
+++ b/test/functional/cases/550_milter_headers.robot
@@ -0,0 +1,39 @@
+*** Settings ***
+Suite Setup     Rspamd Setup
+Suite Teardown  Rspamd Teardown
+Library         ${RSPAMD_TESTDIR}/lib/rspamd.py
+Resource        ${RSPAMD_TESTDIR}/lib/rspamd.robot
+Variables       ${RSPAMD_TESTDIR}/lib/vars.py
+
+*** Variables ***
+${CONFIG}                ${RSPAMD_TESTDIR}/configs/milter_headers.conf
+${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}         {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
new file mode 100644
index 000000000..947bc28dd
--- /dev/null
+++ b/test/functional/configs/milter_headers.conf
@@ -0,0 +1,24 @@
+lua = "{= env.TESTDIR =}/lua/simple.lua"
+
+.include(duplicate=append,priority=0) "{= env.TESTDIR =}/configs/plugins.conf"
+
+milter_headers {
+
+  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 254a6f3ca..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}
@@ -76,10 +98,31 @@ Expect Action
   [Arguments]  ${action}
   Should Be Equal  ${SCAN_RESULT}[action]  ${action}
 
+Expect Added Header
+  [Arguments]  ${header_name}  ${header_value}  ${pos}=-1
+  Dictionary Should Contain Key  ${SCAN_RESULT}  milter
+  ...  msg=milter block was not present in protocol response
+  Dictionary Should Contain Key  ${SCAN_RESULT}[milter]  add_headers
+  ...  msg=add_headers block was not present in protocol response
+  Dictionary Should Contain Key  ${SCAN_RESULT}[milter][add_headers]  ${header_name}
+  ...  msg=${header_name} was not added
+  Should Be Equal  ${SCAN_RESULT}[milter][add_headers][${header_name}][value]  ${header_value}
+  Should Be Equal as Numbers  ${SCAN_RESULT}[milter][add_headers][${header_name}][order]  ${pos}
+
 Expect Email
   [Arguments]  ${email}
   List Should Contain Value  ${SCAN_RESULT}[emails]  ${email}
 
+Expect Removed Header
+  [Arguments]  ${header_name}  ${pos}=0
+  Dictionary Should Contain Key  ${SCAN_RESULT}  milter
+  ...  msg=milter block was not present in protocol response
+  Dictionary Should Contain Key  ${SCAN_RESULT}[milter]  remove_headers
+  ...  msg=remove_headers block was not present in protocol response
+  Dictionary Should Contain Key  ${SCAN_RESULT}[milter][remove_headers]  ${header_name}
+  ...  msg=${header_name} was not removed
+  Should Be Equal as Numbers  ${SCAN_RESULT}[milter][remove_headers][${header_name}]  ${pos}
+
 Expect Required Score
   [Arguments]  ${required_score}
   Should Be Equal As Numbers  ${SCAN_RESULT}[required_score]  ${required_score}