tests/alerts/lib/test_alerttask.py

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'