From eb46eb529a8bcfaeb5a7e736550d7cfb612b02d7 Mon Sep 17 00:00:00 2001
From: GomathiselviS <gomathiselvi@gmail.com>
Date: Mon, 14 Jun 2021 12:42:05 -0400
Subject: Fix keyerror in firewall_rules when state=overridden (#172)

Fix keyerror in firewall_rules when state=overridden

Reviewed-by: https://github.com/apps/ansible-zuul
---
 .../fragments/172-fix-source-keyerror-firewallrules.yml |   3 +++
 .../vyos/config/firewall_rules/firewall_rules.py        |   5 +++--
 .../network/vyos/.test_vyos_firewall_rules.py.swp       | Bin 0 -> 69632 bytes
 .../modules/network/vyos/test_vyos_firewall_rules.py    |  10 ++++++++++
 4 files changed, 16 insertions(+), 2 deletions(-)
 create mode 100644 changelogs/fragments/172-fix-source-keyerror-firewallrules.yml
 create mode 100644 tests/unit/modules/network/vyos/.test_vyos_firewall_rules.py.swp

diff --git a/changelogs/fragments/172-fix-source-keyerror-firewallrules.yml b/changelogs/fragments/172-fix-source-keyerror-firewallrules.yml
new file mode 100644
index 0000000..01825fb
--- /dev/null
+++ b/changelogs/fragments/172-fix-source-keyerror-firewallrules.yml
@@ -0,0 +1,3 @@
+---
+bugfixes:
+  - Fix KeyError 'source' - vyos_firewall_rules
diff --git a/plugins/module_utils/network/vyos/config/firewall_rules/firewall_rules.py b/plugins/module_utils/network/vyos/config/firewall_rules/firewall_rules.py
index ffc642f..fd5a4f5 100644
--- a/plugins/module_utils/network/vyos/config/firewall_rules/firewall_rules.py
+++ b/plugins/module_utils/network/vyos/config/firewall_rules/firewall_rules.py
@@ -779,8 +779,9 @@ class Firewall_rules(ConfigBase):
             key = "group"
             group = w[attr].get(key) or {}
             if group:
-                if h and key in h[attr].keys():
-                    h_group = h[attr].get(key) or {}
+                h_group = {}
+                if h and h.get(attr) and key in h[attr].keys():
+                    h_group = h[attr].get(key)
                 for item, val in iteritems(group):
                     if val:
                         if (
diff --git a/tests/unit/modules/network/vyos/.test_vyos_firewall_rules.py.swp b/tests/unit/modules/network/vyos/.test_vyos_firewall_rules.py.swp
new file mode 100644
index 0000000..b028e81
Binary files /dev/null and b/tests/unit/modules/network/vyos/.test_vyos_firewall_rules.py.swp differ
diff --git a/tests/unit/modules/network/vyos/test_vyos_firewall_rules.py b/tests/unit/modules/network/vyos/test_vyos_firewall_rules.py
index 7219443..682b2da 100644
--- a/tests/unit/modules/network/vyos/test_vyos_firewall_rules.py
+++ b/tests/unit/modules/network/vyos/test_vyos_firewall_rules.py
@@ -963,6 +963,13 @@ class TestVyosFirewallRulesModule(TestVyosModule):
                                         protocol="tcp",
                                         fragment="match-frag",
                                         disabled=False,
+                                        source=dict(
+                                            group=dict(
+                                                address_group="IN-ADDR-GROUP",
+                                                network_group="IN-NET-GROUP",
+                                                port_group="IN-PORT-GROUP",
+                                            )
+                                        ),
                                     ),
                                     dict(
                                         number="2",
@@ -1005,6 +1012,9 @@ class TestVyosFirewallRulesModule(TestVyosModule):
             "set firewall name V4-IN rule 1 protocol 'tcp'",
             "set firewall name V4-IN rule 1 description 'Rule 1 is configured by Ansible RM'",
             "set firewall name V4-IN rule 1 fragment 'match-frag'",
+            "set firewall name V4-IN rule 1 source group address-group IN-ADDR-GROUP",
+            "set firewall name V4-IN rule 1 source group network-group IN-NET-GROUP",
+            "set firewall name V4-IN rule 1 source group port-group IN-PORT-GROUP",
             "set firewall name V4-IN rule 1",
             "set firewall name V4-IN rule 1 action 'reject'",
             "set firewall name V4-IN rule 1 ipsec 'match-ipsec'",
-- 
cgit v1.2.3