[Test] Fix neural plugin tests

3 files changed, 46 insertions, 38 deletions

diff --git a/test/functional/cases/330_neural/001_autotrain.robot b/test/functional/cases/330_neural/001_autotrain.robot
index 4ca29f196..793fc7561 100644
--- a/test/functional/cases/330_neural/001_autotrain.robot
+++ b/test/functional/cases/330_neural/001_autotrain.robot
@@ -16,23 +16,23 @@ ${RSPAMD_URL_TLD}  ${RSPAMD_TESTDIR}/../lua/unit/test_tld.dat
 *** Test Cases ***
 Train
   Sleep  2s  Wait for redis mess
-  FOR    ${INDEX}    IN RANGE    1    11
-    Scan File  ${MESSAGE}  Settings={symbols_enabled = ["SPAM_SYMBOL${INDEX}"]}
+  FOR    ${INDEX}    IN RANGE    4    14
+    Scan File  ${MESSAGE}  Settings={symbols_enabled = ["SPAM_SYMBOL1", "SPAM_SYMBOL2", "SPAM_SYMBOL3", "SPAM_SYMBOL${INDEX}"]}
     Expect Symbol  SPAM_SYMBOL${INDEX}
-    Scan File  ${MESSAGE}  Settings={symbols_enabled = ["HAM_SYMBOL${INDEX}"]}
+    Scan File  ${MESSAGE}  Settings={symbols_enabled = ["HAM_SYMBOL1", "HAM_SYMBOL2", "HAM_SYMBOL3", "HAM_SYMBOL${INDEX}"]}
     Expect Symbol  HAM_SYMBOL${INDEX}
   END
 
 Check Neural HAM
-  Sleep  2s  Wait for neural to be loaded
-  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["HAM_SYMBOL1","HAM_SYMBOL2","HAM_SYMBOL5"];groups_enabled=["neural"];symbols_disabled = ["NEURAL_LEARN"]}
+  Sleep  5s  Wait for neural to be loaded
+  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["HAM_SYMBOL1","HAM_SYMBOL2","HAM_SYMBOL3","HAM_SYMBOL8"];groups_enabled=["neural"];symbols_disabled = ["NEURAL_LEARN"]}
   Expect Symbol  NEURAL_HAM_SHORT
   Do Not Expect Symbol  NEURAL_SPAM_SHORT
   #Expect Symbol  NEURAL_HAM_SHORT_PCA
   #Do Not Expect Symbol  NEURAL_SPAM_SHORT_PCA
 
 Check Neural SPAM
-  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["SPAM_SYMBOL2","SPAM_SYMBOL4","SPAM_SYMBOL6","SPAM_SYMBOL8"];groups_enabled=["neural"];symbols_disabled = ["NEURAL_LEARN"]}
+  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["SPAM_SYMBOL1","SPAM_SYMBOL2","SPAM_SYMBOL3","SPAM_SYMBOL8"];groups_enabled=["neural"];symbols_disabled = ["NEURAL_LEARN"]}
   Expect Symbol  NEURAL_SPAM_SHORT
   Do Not Expect Symbol  NEURAL_HAM_SHORT
   #Expect Symbol  NEURAL_SPAM_SHORT_PCA
@@ -40,24 +40,24 @@ Check Neural SPAM
 
 
 Train INVERSE
-  FOR    ${INDEX}    IN RANGE    1    11
-    Scan File  ${MESSAGE}  Settings={symbols_enabled = ["SPAM_SYMBOL${INDEX}"]; SPAM_SYMBOL${INDEX} = -5;}
+  FOR    ${INDEX}    IN RANGE    4    14
+    Scan File  ${MESSAGE}  Settings={symbols_enabled = ["SPAM_SYMBOL1", "SPAM_SYMBOL2", "SPAM_SYMBOL3","SPAM_SYMBOL${INDEX}"]; SPAM_SYMBOL1 = -5; SPAM_SYMBOL2 = -5; SPAM_SYMBOL3 = -5; SPAM_SYMBOL${INDEX} = -5;}
     Expect Symbol  SPAM_SYMBOL${INDEX}
-    Scan File  ${MESSAGE}  Settings={symbols_enabled = ["HAM_SYMBOL${INDEX}"]; HAM_SYMBOL${INDEX} = 5;}
+    Scan File  ${MESSAGE}  Settings={symbols_enabled = ["HAM_SYMBOL1", "HAM_SYMBOL2", "HAM_SYMBOL3","HAM_SYMBOL${INDEX}"]; HAM_SYMBOL1 = 5; HAM_SYMBOL2 = 5; HAM_SYMBOL3 = 5; HAM_SYMBOL${INDEX} = 5;}
     Expect Symbol  HAM_SYMBOL${INDEX}
   END
 
 Check Neural HAM INVERSE
-  Sleep  2s  Wait for neural to be loaded
-  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["HAM_SYMBOL1","HAM_SYMBOL2","HAM_SYMBOL5"];groups_enabled=["neural"]}
+  Sleep  5s  Wait for neural to be loaded
+  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["HAM_SYMBOL1","HAM_SYMBOL2","HAM_SYMBOL3","HAM_SYMBOL5",];groups_enabled=["neural"]}
   Expect Symbol  NEURAL_SPAM_SHORT
   #Expect Symbol  NEURAL_SPAM_SHORT_PCA
   Do Not Expect Symbol  NEURAL_HAM_SHORT
   #Do Not Expect Symbol  NEURAL_HAM_SHORT_PCA
 
 Check Neural SPAM INVERSE
-  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["SPAM_SYMBOL2","SPAM_SYMBOL4","SPAM_SYMBOL6","SPAM_SYMBOL8"];groups_enabled=["neural"]}
+  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["SPAM_SYMBOL1","SPAM_SYMBOL2","SPAM_SYMBOL3","SPAM_SYMBOL5"];groups_enabled=["neural"]}
   Expect Symbol  NEURAL_HAM_SHORT
   #Expect Symbol  NEURAL_HAM_SHORT_PCA
   Do Not Expect Symbol  NEURAL_SPAM_SHORT
-  #Do Not Expect Symbol  NEURAL_SPAM_SHORT_PCA
+  #Do Not Expect Symbol  NEURAL_SPAM_SHORT_PCA
\ No newline at end of file
diff --git a/test/functional/cases/330_neural/002_manualtrain.robot b/test/functional/cases/330_neural/002_manualtrain.robot
index ecee7f631..bbc6c2e17 100644
--- a/test/functional/cases/330_neural/002_manualtrain.robot
+++ b/test/functional/cases/330_neural/002_manualtrain.robot
@@ -15,50 +15,58 @@ ${RSPAMD_URL_TLD}  ${RSPAMD_TESTDIR}/../lua/unit/test_tld.dat
 
 *** Test Cases ***
 Collect training vectors & train manually
-  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["SPAM_SYMBOL","SAVE_NN_ROW"]}
-  Expect Symbol  SPAM_SYMBOL
-  # Save neural inputs for later
-  ${SPAM_ROW} =  Get File  ${SCAN_RESULT}[symbols][SAVE_NN_ROW][options][0]
-  Remove File  ${SCAN_RESULT}[symbols][SAVE_NN_ROW][options][0]
-  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["HAM_SYMBOL","SAVE_NN_ROW"]}
-  Expect Symbol  HAM_SYMBOL
-  # Save neural inputs for later
-  ${HAM_ROW} =  Get File  ${SCAN_RESULT}[symbols][SAVE_NN_ROW][options][0]
-  Remove File  ${SCAN_RESULT}[symbols][SAVE_NN_ROW][options][0]
-  ${HAM_ROW} =  Run  ${RSPAMADM} lua -a ${HAM_ROW} ${RSPAMD_TESTDIR}/util/nn_unpack.lua
-  ${HAM_ROW} =  Evaluate  json.loads("${HAM_ROW}")
-  ${SPAM_ROW} =  Run  ${RSPAMADM} lua -a ${SPAM_ROW} ${RSPAMD_TESTDIR}/util/nn_unpack.lua
-  ${SPAM_ROW} =  Evaluate  json.loads("${SPAM_ROW}")
-  ${HAM_VEC} =  Evaluate  [${HAM_ROW}] * 10
-  ${SPAM_VEC} =  Evaluate  [${SPAM_ROW}] * 10
+  @{HAM_VEC}=    Create List
+  @{SPAM_VEC}=    Create List
+
+  FOR    ${INDEX}    IN RANGE    4    14
+    Scan File  ${MESSAGE}  Settings={symbols_enabled = ["SPAM_SYMBOL1", "SPAM_SYMBOL2", "SPAM_SYMBOL3", "SPAM_SYMBOL${INDEX}", "SAVE_NN_ROW"]}
+    Expect Symbol  SPAM_SYMBOL${INDEX}
+    # Save neural inputs for later
+    ${SPAM_ROW} =  Get File  ${SCAN_RESULT}[symbols][SAVE_NN_ROW][options][0]
+    Remove File  ${SCAN_RESULT}[symbols][SAVE_NN_ROW][options][0]
+    Scan File  ${MESSAGE}  Settings={symbols_enabled = ["HAM_SYMBOL1", "HAM_SYMBOL2", "HAM_SYMBOL3", "HAM_SYMBOL${INDEX}", "SAVE_NN_ROW"]}
+    Expect Symbol  HAM_SYMBOL${INDEX}
+    # Save neural inputs for later
+    ${HAM_ROW} =  Get File  ${SCAN_RESULT}[symbols][SAVE_NN_ROW][options][0]
+    Remove File  ${SCAN_RESULT}[symbols][SAVE_NN_ROW][options][0]
+    ${HAM_ROW} =  Run  ${RSPAMADM} lua -a ${HAM_ROW} ${RSPAMD_TESTDIR}/util/nn_unpack.lua
+    ${HAM_ROW} =  Evaluate  json.loads("${HAM_ROW}")
+    ${SPAM_ROW} =  Run  ${RSPAMADM} lua -a ${SPAM_ROW} ${RSPAMD_TESTDIR}/util/nn_unpack.lua
+    ${SPAM_ROW} =  Evaluate  json.loads("${SPAM_ROW}")
+    Append To List    ${HAM_VEC}    ${HAM_ROW}
+    Append To List    ${SPAM_VEC}   ${SPAM_ROW}
+  END
+
   ${json1} =  Evaluate  json.dumps({"spam_vec": ${SPAM_VEC}, "ham_vec": ${HAM_VEC}, "rule": "SHORT"})
+  HTTP  POST  ${RSPAMD_LOCAL_ADDR}  ${RSPAMD_PORT_CONTROLLER}  /plugins/neural/learn  ${json1}
+
   # Save variables for use in inverse training
   Set Suite Variable  ${HAM_VEC}
   Set Suite Variable  ${SPAM_VEC}
-  HTTP  POST  ${RSPAMD_LOCAL_ADDR}  ${RSPAMD_PORT_CONTROLLER}  /plugins/neural/learn  ${json1}
-  Sleep  2s  Wait for neural to be loaded
+
+  Sleep  5s  Wait for neural to be loaded
 
 Check Neural HAM
-  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["HAM_SYMBOL"];groups_enabled=["neural"];symbols_disabled = ["NEURAL_LEARN"]}
+  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["HAM_SYMBOL1","HAM_SYMBOL2","HAM_SYMBOL3","HAM_SYMBOL8"];groups_enabled=["neural"];symbols_disabled = ["NEURAL_LEARN"]}
   Do Not Expect Symbol  NEURAL_SPAM_SHORT
   Expect Symbol  NEURAL_HAM_SHORT
 
 Check Neural SPAM
-  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["SPAM_SYMBOL"];groups_enabled=["neural"];symbols_disabled = ["NEURAL_LEARN"]}
+  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["SPAM_SYMBOL1","SPAM_SYMBOL2","SPAM_SYMBOL3","SPAM_SYMBOL8"];groups_enabled=["neural"];symbols_disabled = ["NEURAL_LEARN"]}
   Do Not Expect Symbol  NEURAL_HAM_SHORT
   Expect Symbol  NEURAL_SPAM_SHORT
 
 Train inverse
   ${json2} =  Evaluate  json.dumps({"spam_vec": ${HAM_VEC}, "ham_vec": ${SPAM_VEC}, "rule": "SHORT"})
   HTTP  POST  ${RSPAMD_LOCAL_ADDR}  ${RSPAMD_PORT_CONTROLLER}  /plugins/neural/learn  ${json2}
-  Sleep  2s  Wait for neural to be loaded
+  Sleep  5s  Wait for neural to be loaded
 
 Check Neural HAM - inverse
-  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["HAM_SYMBOL"];groups_enabled=["neural"];symbols_disabled = ["NEURAL_LEARN"]}
+  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["HAM_SYMBOL1","HAM_SYMBOL2","HAM_SYMBOL3","HAM_SYMBOL8"];groups_enabled=["neural"];symbols_disabled = ["NEURAL_LEARN"]}
   Do Not Expect Symbol  NEURAL_HAM_SHORT
   Expect Symbol  NEURAL_SPAM_SHORT
 
 Check Neural SPAM - inverse
-  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["SPAM_SYMBOL"];groups_enabled=["neural"];symbols_disabled = ["NEURAL_LEARN"]}
+  Scan File  ${MESSAGE}  Settings={symbols_enabled = ["SPAM_SYMBOL1","SPAM_SYMBOL2","SPAM_SYMBOL3","SPAM_SYMBOL8"];groups_enabled=["neural"];symbols_disabled = ["NEURAL_LEARN"]}
   Do Not Expect Symbol  NEURAL_SPAM_SHORT
-  Expect Symbol  NEURAL_HAM_SHORT
+  Expect Symbol  NEURAL_HAM_SHORT
\ No newline at end of file
diff --git a/test/functional/lua/neural.lua b/test/functional/lua/neural.lua
index 32419c8c8..fb599f3b0 100644
--- a/test/functional/lua/neural.lua
+++ b/test/functional/lua/neural.lua
@@ -1,6 +1,6 @@
 local logger = require "rspamd_logger"
 
-for i = 1,10 do
+for i = 1,14 do
   rspamd_config:register_symbol({
     name = 'SPAM_SYMBOL'..tostring(i),
     score = 5.0,