webssh/tests/test_policy.py

import os
import unittest
import paramiko

from paramiko.client import RejectPolicy, WarningPolicy
from policy import (AutoAddPolicy, get_policy_dictionary, load_host_keys,
                    get_policy_class, check_policy_setting)


class TestPolicy(unittest.TestCase):

    def test_get_policy_dictionary(self):
        classes = [AutoAddPolicy, RejectPolicy, WarningPolicy]
        dic = get_policy_dictionary()
        for cls in classes:
            val = dic[cls.__name__.lower()]
            self.assertIs(cls, val)

    def test_load_host_keys(self):
        path = '/path-not-exists'
        host_keys = load_host_keys(path)
        self.assertFalse(host_keys)

        path = '/tmp'
        host_keys = load_host_keys(path)
        self.assertFalse(host_keys)

        path = 'tests/known_hosts_example'
        host_keys = load_host_keys(path)
        self.assertEqual(host_keys, paramiko.hostkeys.HostKeys(path))

    def test_get_policy_class(self):
        keys = ['autoadd', 'reject', 'warning']
        vals = [AutoAddPolicy, RejectPolicy, WarningPolicy]
        for key, val in zip(keys, vals):
            cls = get_policy_class(key)
            self.assertIs(cls, val)

        key = 'non-exists'
        with self.assertRaises(ValueError):
            get_policy_class(key)

    def test_check_policy_setting(self):
        host_keys_filename = './tests/host_keys_test.db'
        host_keys_settings = dict(
            host_keys=paramiko.hostkeys.HostKeys(),
            system_host_keys=paramiko.hostkeys.HostKeys(),
            host_keys_filename=host_keys_filename
        )

        with self.assertRaises(ValueError):
            check_policy_setting(RejectPolicy, host_keys_settings)

        try:
            os.unlink(host_keys_filename)
        except OSError:
            pass
        check_policy_setting(AutoAddPolicy, host_keys_settings)
        self.assertEqual(os.path.exists(host_keys_filename), True)