ipfw: add state/comment tests

1 files changed, 61 insertions, 2 deletions

diff --git a/sbin/ipfw/tests/test_add_rule.py b/sbin/ipfw/tests/test_add_rule.py
index 42c594a83bcf..60c8cebaceaa 100755
--- a/sbin/ipfw/tests/test_add_rule.py
+++ b/sbin/ipfw/tests/test_add_rule.py
@@ -57,7 +57,7 @@ def differ(w_obj, g_obj, w_stack=[], g_stack=[]):
         return True
     num_objects = 0
     for i, w_child in enumerate(w_obj.obj_list):
-        if i > len(g_obj.obj_list):
+        if i >= len(g_obj.obj_list):
             print("MISSING object from chain {}".format(" / ".join(w_stack)))
             w_child.print_obj()
             print("==========================")
@@ -206,6 +206,66 @@ class TestAddRule(BaseTest):
                 },
                 id="test_eaction_ntp",
             ),
+            pytest.param(
+                {
+                    "in": "add // test comment",
+                    "out": {
+                        "insns": [
+                            InsnComment(comment="test comment"),
+                            Insn(IpFwOpcode.O_COUNT),
+                        ],
+                    },
+                },
+                id="test_action_comment",
+            ),
+            pytest.param(
+                {
+                    "in": "add check-state :OUT // test comment",
+                    "out": {
+                        "objs": [
+                            NTlv(IpFwTlvType.IPFW_TLV_STATE_NAME, idx=1, name="OUT"),
+                        ],
+                        "insns": [
+                            InsnComment(comment="test comment"),
+                            Insn(IpFwOpcode.O_CHECK_STATE, arg1=1),
+                        ],
+                    },
+                },
+                id="test_check_state",
+            ),
+            pytest.param(
+                {
+                    "in": "add allow tcp from any to any keep-state :OUT",
+                    "out": {
+                        "objs": [
+                            NTlv(IpFwTlvType.IPFW_TLV_STATE_NAME, idx=1, name="OUT"),
+                        ],
+                        "insns": [
+                            Insn(IpFwOpcode.O_PROBE_STATE, arg1=1),
+                            Insn(IpFwOpcode.O_PROTO, arg1=6),
+                            Insn(IpFwOpcode.O_KEEP_STATE, arg1=1),
+                            InsnEmpty(IpFwOpcode.O_ACCEPT),
+                        ],
+                    },
+                },
+                id="test_keep_state",
+            ),
+            pytest.param(
+                {
+                    "in": "add allow tcp from any to any record-state",
+                    "out": {
+                        "objs": [
+                            NTlv(IpFwTlvType.IPFW_TLV_STATE_NAME, idx=1, name="default"),
+                        ],
+                        "insns": [
+                            Insn(IpFwOpcode.O_PROTO, arg1=6),
+                            Insn(IpFwOpcode.O_KEEP_STATE, arg1=1),
+                            InsnEmpty(IpFwOpcode.O_ACCEPT),
+                        ],
+                    },
+                },
+                id="test_record_state",
+            ),
         ],
     )
     def test_add_rule(self, rule):
@@ -329,7 +389,6 @@ class TestAddRule(BaseTest):
                 ("call 420", Insn(IpFwOpcode.O_CALLRETURN, arg1=420)), id="call_420"
             ),
             # TOK_FORWARD
-            # TOK_COMMENT
             pytest.param(
                 ("setfib 1", Insn(IpFwOpcode.O_SETFIB, arg1=1 | 0x8000)),
                 id="setfib_1",