update_connected_ips: correctly handle byte-string

This commit is contained in:
Pawel Marczewski 2020-01-14 10:14:00 +01:00
parent 39885a4329
commit e43fd2fc5a
No known key found for this signature in database
GPG Key ID: DE42EE9B14F96465
2 changed files with 8 additions and 12 deletions

View File

@ -81,10 +81,8 @@ class FirewallWorker(object):
raise NotImplementedError raise NotImplementedError
def get_connected_ips(self, family): def get_connected_ips(self, family):
if family == 6: ips = self.qdb.read('/connected-ips6' if family == 6 else '/connected-ips')
return self.qdb.read('/connected-ips6').split() return ips.decode().split()
else:
return self.qdb.read('/connected-ips').split()
def run_firewall_dir(self): def run_firewall_dir(self):
"""Run scripts dir contents, before user script""" """Run scripts dir contents, before user script"""

View File

@ -332,8 +332,7 @@ class TestIptablesWorker(TestCase):
]) ])
def test_008_update_connected_ips(self): def test_008_update_connected_ips(self):
with patch.object(self.obj, 'get_connected_ips') as get_connected_ips: self.obj.qdb.entries['/connected-ips'] = b'10.137.0.1 10.137.0.2'
get_connected_ips.return_value = ['10.137.0.1', '10.137.0.2']
self.obj.called_commands[4] = [] self.obj.called_commands[4] = []
self.obj.update_connected_ips(4) self.obj.update_connected_ips(4)
@ -529,8 +528,7 @@ class TestNftablesWorker(TestCase):
]) ])
def test_008_update_connected_ips(self): def test_008_update_connected_ips(self):
with patch.object(self.obj, 'get_connected_ips') as get_connected_ips: self.obj.qdb.entries['/connected-ips'] = b'10.137.0.1 10.137.0.2'
get_connected_ips.return_value = ['10.137.0.1', '10.137.0.2']
self.obj.loaded_rules = [] self.obj.loaded_rules = []
self.obj.update_connected_ips(4) self.obj.update_connected_ips(4)