Fix migration tests

skalenetwork · Dec 23, 2024 · 35e33aa · 35e33aa
1 parent cd0bce6
commit 35e33aa
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/tests/core/migration_test.py b/tests/core/migration_test.py
@@ -4,8 +4,11 @@
 
 from node_cli.utils.helper import run_cmd
 
+CUSTOM_CHAIN_NAME = 'TEST'
+
 
 def add_base_rules():
+    run_cmd(f'iptables -N {CUSTOM_CHAIN_NAME}'.split(' '))
     run_cmd('iptables -A INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT'.split(' '))
     run_cmd('iptables -A INPUT -i lo -j ACCEPT'.split(' '))
     run_cmd('iptables -A INPUT -p tcp --dport 22 -j ACCEPT'.split(' '))
@@ -15,7 +18,7 @@ def add_base_rules():
     run_cmd('iptables -A INPUT -p udp --dport 53 -j ACCEPT'.split(' '))
     run_cmd('iptables -A INPUT -p tcp --dport 3009 -j ACCEPT'.split(' '))
     # non skale related rule
-    run_cmd('iptables -A INPUT -p tcp --dport 2222 -j ACCEPT'.split(' '))
+    run_cmd(f'iptables -A {CUSTOM_CHAIN_NAME} -p tcp --dport 2222 -j ACCEPT'.split(' '))
     run_cmd('iptables -A INPUT -p tcp --dport 9100 -j ACCEPT'.split(' '))
     run_cmd('iptables -A INPUT -p tcp -j DROP'.split(' '))
     run_cmd('iptables -A INPUT -p udp -j DROP'.split(' '))
@@ -37,4 +40,5 @@ def test_migration(base_rules):
     migrate()
     res = run_cmd(['iptables', '-S'])
     output = res.stdout.decode('utf-8')
-    assert output == '-P INPUT ACCEPT\n-P FORWARD ACCEPT\n-P OUTPUT ACCEPT\n-A INPUT -p tcp -m tcp --dport 2222 -j ACCEPT\n'  # noqa
+    print(output)
+    assert output == f'-P INPUT ACCEPT\n-P FORWARD ACCEPT\n-P OUTPUT ACCEPT\n-N {CUSTOM_CHAIN_NAME}\n-A {CUSTOM_CHAIN_NAME} -p tcp -m tcp --dport 2222 -j ACCEPT\n'  # noqa