File size: 1,964 Bytes
7c89ed7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | import mock
import socket
import os
import sys
def reverse_lookup(ip):
if ip == '10.1.1.1':
return ('test.domain.com', ip)
if ip == '8.8.8.8':
raise socket.herror
class AlertTaskTest(object):
def teardown(self):
sys.path.remove(self.alerts_path)
sys.path.remove(self.alerts_lib_path)
if 'lib' in sys.modules:
del sys.modules['lib']
def setup(self):
self.alerts_path = os.path.join(os.path.dirname(__file__), "../../../alerts")
self.alerts_lib_path = os.path.join(os.path.dirname(__file__), "../../../alerts/lib")
sys.path.insert(0, self.alerts_path)
sys.path.insert(1, self.alerts_lib_path)
from lib import alerttask
self.hostname_from_ip = alerttask.hostname_from_ip
self.add_hostname_to_ip = alerttask.add_hostname_to_ip
class TestHostnameFromIP(AlertTaskTest):
def test_good_hostname_from_ip(self):
with mock.patch("socket.gethostbyaddr", side_effect=reverse_lookup):
hostname = self.hostname_from_ip('10.1.1.1')
assert hostname == 'test.domain.com'
def test_bad_hostname_from_ip(self):
with mock.patch("socket.gethostbyaddr", side_effect=reverse_lookup):
hostname = self.hostname_from_ip('8.8.8.8')
assert hostname is None
class TestAddHostnameToIP(AlertTaskTest):
def setup(self):
super().setup()
self.formatted_string = '{0} ({1})'
def test_internal_hostname(self):
with mock.patch("socket.gethostbyaddr", side_effect=reverse_lookup):
hostname_info = self.add_hostname_to_ip('10.1.1.1', self.formatted_string)
assert hostname_info == '10.1.1.1 (test.domain.com)'
def test_external_hostname(self):
with mock.patch("socket.gethostbyaddr", side_effect=reverse_lookup):
hostname_info = self.add_hostname_to_ip('8.8.8.8', self.formatted_string)
assert hostname_info == '8.8.8.8'
|