[Test] Add test for TLD parts

5 files changed, 72 insertions, 1 deletions

diff --git a/test/functional/cases/101_lua.robot b/test/functional/cases/101_lua.robot
index c465491f1..27bb867f7 100644
--- a/test/functional/cases/101_lua.robot
+++ b/test/functional/cases/101_lua.robot
@@ -8,6 +8,7 @@ Variables       ${TESTDIR}/lib/vars.py
 ${CONFIG}       ${TESTDIR}/configs/lua_test.conf
 ${MESSAGE}      ${TESTDIR}/messages/spam_message.eml
 ${RSPAMD_SCOPE}  Test
+${URL_TLD}      ${TESTDIR}/../lua/unit/test_tld.dat
 
 *** Test Cases ***
 Flags
@@ -34,8 +35,18 @@ Recipient Parsing Sanity
   ...  -r  rcpt3@foobar  -r  rcpt4@foobar
   Check Rspamc  ${result}  TEST_RCPT (1.00)[rcpt1@foobar,rcpt2@foobar,rcpt3@foobar,rcpt4@foobar]
 
+TLD parts
+  [Setup]  TLD Setup  ${TESTDIR}/lua/tlds.lua
+  ${result} =  Scan Message With Rspamc  ${MESSAGE}
+  Check Rspamc  ${result}  TEST_TLD (1.00)[no worry]
+
 *** Keywords ***
 Lua Setup
   [Arguments]  ${LUA_SCRIPT}
   Set Test Variable  ${LUA_SCRIPT}
   Generic Setup
+
+TLD Setup
+  [Arguments]  ${LUA_SCRIPT}
+  Set Test Variable  ${URL_TLD}  ${TESTDIR}/../../contrib/publicsuffix/effective_tld_names.dat
+  Lua Setup  ${LUA_SCRIPT}
diff --git a/test/functional/cases/121_json/lib.robot b/test/functional/cases/121_json/lib.robot
index e558c7107..7729bcc3a 100644
--- a/test/functional/cases/121_json/lib.robot
+++ b/test/functional/cases/121_json/lib.robot
@@ -1,3 +1,6 @@
+*** Variables ***
+${URL_TLD}      ${TESTDIR}/../lua/unit/test_tld.dat
+
 *** Keywords ***
 Stat Test
   @{result} =  HTTP  GET  ${LOCAL_ADDR}  ${PORT_CONTROLLER}  /stat
diff --git a/test/functional/cases/140_proxy.robot b/test/functional/cases/140_proxy.robot
index a82c687e4..3907cfed6 100644
--- a/test/functional/cases/140_proxy.robot
+++ b/test/functional/cases/140_proxy.robot
@@ -8,6 +8,7 @@ Variables       ${TESTDIR}/lib/vars.py
 *** Variables ***
 ${LUA_SCRIPT}   ${TESTDIR}/lua/simple.lua
 ${MESSAGE}      ${TESTDIR}/messages/spam_message.eml
+${URL_TLD}      ${TESTDIR}/../lua/unit/test_tld.dat
 
 *** Test Cases ***
 Scan Message
diff --git a/test/functional/configs/lua_test.conf b/test/functional/configs/lua_test.conf
index bcb1c3323..6e93f85ac 100644
--- a/test/functional/configs/lua_test.conf
+++ b/test/functional/configs/lua_test.conf
@@ -1,6 +1,6 @@
 options = {
 	filters = ["spf", "dkim", "regexp"]
-	url_tld = "${TESTDIR}/../lua/unit/test_tld.dat"
+	url_tld = "${URL_TLD}"
 	pidfile = "${TMPDIR}/rspamd.pid"
 }
 logging = {
diff --git a/test/functional/lua/tlds.lua b/test/functional/lua/tlds.lua
new file mode 100644
index 000000000..9f2d64feb
--- /dev/null
+++ b/test/functional/lua/tlds.lua
@@ -0,0 +1,56 @@
+rspamd_config:register_symbol({
+  name = 'TEST_TLD',
+  score = 1.0,
+  callback = function(task)
+    local prefixes = {
+      '',
+      'example.'
+    }
+    local test_domains = {
+      'example.ac',
+      'example.b.br',
+      'example.co',
+      'example.city.kawasaki.jp',
+      'example.com',
+      'example.co.za',
+      'example.in.net',
+      'example.kawasaki.jp',
+      'example.net',
+      'example.net.in',
+      'example.nom.br',
+      'example.org',
+      'example.org.ac',
+      'example.ru.com',
+      'example.za.net',
+      'example.za.org',
+      'org.org.za',
+    }
+    local worry = {}
+    local rspamd_mempool = require 'rspamd_mempool'
+    local rspamd_url = require 'rspamd_url'
+    local rspamd_util = require 'rspamd_util'
+    local pool = rspamd_mempool.create()
+    for _, d in ipairs(test_domains) do
+      (function()
+        for _, p in ipairs(prefixes) do
+          local test = rspamd_util.get_tld(p .. d)
+          if (test ~= d) then
+            table.insert(worry, 'util.get_tld:' .. p .. d .. ':' .. test)
+            return
+          end
+          local u = rspamd_url.create(pool, p .. d)
+          local test = u:get_tld()
+          if (test ~= d) then
+            table.insert(worry, 'url.get_tld:' .. p .. d .. ':' .. test)
+            return
+          end
+        end
+      end)()
+    end
+    if (#worry == 0) then
+      return true, "no worry"
+    else
+      return true, table.concat(worry, ",")
+    end
+  end
+})