1""" (Very) basic ip4trie dataset tests
2"""
3import unittest
4
5from rbldnsd import Rbldnsd, ZoneFile
6
7__all__ = [
8    'TestIp4TrieDataset',
9    ]
10
11def ip4trie(zone_data):
12    """ Run rbldnsd with an ip6trie dataset
13    """
14    dnsd = Rbldnsd()
15    dnsd.add_dataset('ip4trie', ZoneFile(zone_data))
16    return dnsd
17
18def reversed_ip(ip4addr, domain='example.com'):
19    revip = '.'.join(reversed(ip4addr.split('.')))
20    return "%s.%s" % (revip, domain)
21
22class TestIp4TrieDataset(unittest.TestCase):
23    def test_exclusion(self):
24        with ip4trie(["1.2.3.0/24 listed",
25                      "!1.2.3.4"]) as dnsd:
26            self.assertEqual(dnsd.query(reversed_ip("1.2.3.4")), None)
27            self.assertEqual(dnsd.query(reversed_ip("1.2.3.3")), "listed")
28            self.assertEqual(dnsd.query(reversed_ip("1.2.3.5")), "listed")
29
30    def test_wildcard_prefix(self):
31        with ip4trie(["0/0 wild",
32                      "127.0.0.1 localhost"]) as dnsd:
33            self.assertEqual(dnsd.query(reversed_ip("127.0.0.1")), "localhost")
34            self.assertEqual(dnsd.query(reversed_ip("0.0.0.0")), "wild")
35            self.assertEqual(dnsd.query(reversed_ip("127.0.0.2")), "wild")
36
37if __name__ == '__main__':
38    unittest.main()
39