From 2adcd0a6f7318d62871727c866aad5640e552e70 Mon Sep 17 00:00:00 2001
From: Andrew Lewis <nerf@judo.za.org>
Date: Tue, 30 Jul 2024 14:37:13 +0200
Subject: [PATCH] [Test] Add tests for milter_headers

---
 .../functional/cases/550_milter_headers.robot | 20 ++++++++++++++++++
 test/functional/configs/milter_headers.conf   | 15 +++++++++++++
 test/functional/lib/rspamd.robot              | 21 +++++++++++++++++++
 3 files changed, 56 insertions(+)
 create mode 100644 test/functional/cases/550_milter_headers.robot
 create mode 100644 test/functional/configs/milter_headers.conf

diff --git a/test/functional/cases/550_milter_headers.robot b/test/functional/cases/550_milter_headers.robot
new file mode 100644
index 000000000..737d0b42d
--- /dev/null
+++ b/test/functional/cases/550_milter_headers.robot
@@ -0,0 +1,20 @@
+*** 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}         {symbols_enabled = [SIMPLE_TEST]}
+
+*** Test Cases ***
+CHECK HEADERS WITH TEST SYMBOL
+  Scan File  ${MESSAGE}  Settings=${SETTINGS_TEST}
+  Expect Removed Header  X-Virus
+  Expect Added Header  X-Virus  Fires always
diff --git a/test/functional/configs/milter_headers.conf b/test/functional/configs/milter_headers.conf
new file mode 100644
index 000000000..0787e9463
--- /dev/null
+++ b/test/functional/configs/milter_headers.conf
@@ -0,0 +1,15 @@
+lua = "{= env.TESTDIR =}/lua/simple.lua"
+
+.include(duplicate=append,priority=0) "{= env.TESTDIR =}/configs/plugins.conf"
+
+milter_headers {
+
+  use = ["x-virus"];
+
+  routines {
+    x-virus {
+      symbols = ["SIMPLE_TEST"];
+    }
+  }
+
+}
diff --git a/test/functional/lib/rspamd.robot b/test/functional/lib/rspamd.robot
index 254a6f3ca..917b01298 100644
--- a/test/functional/lib/rspamd.robot
+++ b/test/functional/lib/rspamd.robot
@@ -76,10 +76,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}
-- 
2.39.5