Run tests with both ipaddress object and string IPs

Author: oschwald

tests/reader_test.py

@@ -4,6 +4,7 @@
 from __future__ import unicode_literals
 
 import logging
+import ipaddress
 import mock
 import os
 import sys
@@ -48,6 +49,13 @@ def get_reader_from_file_descriptor(filepath, mode):
 
 
 class BaseTestReader(object):
+    use_ip_objects = False
+
+    def ipf(self, ip):
+        if self.use_ip_objects:
+            return ipaddress.ip_address(ip)
+        return ip
+
     def test_reader(self):
         for record_size in [24, 28, 32]:
             for ip_version in [4, 6]:
@@ -67,7 +75,7 @@ def test_reader(self):
     def test_decoder(self):
         reader = open_database(
             'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
-        record = reader.get('::1.1.1.0')
+        record = reader.get(self.ipf('::1.1.1.0'))
 
         self.assertEqual(record['array'], [1, 2, 3])
         self.assertEqual(record['boolean'], True)
@@ -97,8 +105,8 @@ def test_no_ipv4_search_tree(self):
             'tests/data/test-data/MaxMind-DB-no-ipv4-search-tree.mmdb',
             self.mode)
 
-        self.assertEqual(reader.get('1.1.1.1'), '::0/64')
-        self.assertEqual(reader.get('192.1.1.1'), '::0/64')
+        self.assertEqual(reader.get(self.ipf('1.1.1.1')), '::0/64')
+        self.assertEqual(reader.get(self.ipf('192.1.1.1')), '::0/64')
         reader.close()
 
     def test_ipv6_address_in_ipv4_database(self):
@@ -108,7 +116,7 @@ def test_ipv6_address_in_ipv4_database(self):
                 ValueError, 'Error looking up 2001::. '
                 'You attempted to look up an IPv6 address '
                 'in an IPv4-only database'):
-            reader.get('2001::')
+            reader.get(self.ipf('2001::'))
         reader.close()
 
     def test_no_extension_exception(self):
@@ -122,14 +130,6 @@ def test_no_extension_exception(self):
                           MODE_MMAP_EXT)
         maxminddb.extension = real_extension
 
-    def test_ip_object_lookup(self):
-        reader = open_database('tests/data/test-data/GeoIP2-City-Test.mmdb',
-                               self.mode)
-        with self.assertRaisesRegex(TypeError,
-                                    "must be str(?:ing)?, not IPv6Address"):
-            reader.get(compat_ip_address('2001:220::'))
-        reader.close()
-
     def test_broken_database(self):
         reader = open_database(
             'tests/data/test-data/'
@@ -138,7 +138,7 @@ def test_broken_database(self):
                 InvalidDatabaseError, r"The MaxMind DB file's data "
                 r"section contains bad data \(unknown data "
                 r"type or corrupt data\)"):
-            reader.get('2001:220::')
+            reader.get(self.ipf('2001:220::'))
         reader.close()
 
     def test_ip_validation(self):
@@ -183,7 +183,7 @@ def test_too_many_get_args(self):
         reader = open_database(
             'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
         with self.assertRaises(TypeError):
-            reader.get('1.1.1.1', 'blah')
+            reader.get(self.ipf('1.1.1.1'), 'blah')
         reader.close()
 
     def test_no_get_args(self):
@@ -231,7 +231,7 @@ def test_closed_get(self):
         with self.assertRaisesRegex(
                 ValueError,
                 'Attempt to read from a closed MaxMind DB.|closed'):
-            reader.get('1.1.1.1')
+            reader.get(self.ipf('1.1.1.1'))
 
     def test_with_statement(self):
         filename = 'tests/data/test-data/MaxMind-DB-test-ipv4-24.mmdb'
@@ -296,7 +296,7 @@ def _check_concurrency(self, worker_class):
         def lookup(pipe):
             try:
                 for i in range(32):
-                    reader.get('65.115.240.{i}'.format(i=i))
+                    reader.get(self.ipf('65.115.240.{i}'.format(i=i)))
                 pipe.send(1)
             except:
                 pipe.send(0)
@@ -340,7 +340,7 @@ def _check_metadata(self, reader, ip_version, record_size):
     def _check_ip_v4(self, reader, file_name):
         for i in range(6):
             address = '1.1.1.' + str(pow(2, i))
-            self.assertEqual({'ip': address}, reader.get(address),
+            self.assertEqual({'ip': address}, reader.get(self.ipf(address)),
                              'found expected data record for ' + address +
                              ' in ' + file_name)
 
@@ -357,20 +357,20 @@ def _check_ip_v4(self, reader, file_name):
             data = {'ip': value_address}
 
             self.assertEqual(
-                data, reader.get(key_address),
+                data, reader.get(self.ipf(key_address)),
                 'found expected data record for ' + key_address + ' in ' +
                 file_name)
 
         for ip in ['1.1.1.33', '255.254.253.123']:
-            self.assertIsNone(reader.get(ip))
+            self.assertIsNone(reader.get(self.ipf(ip)))
 
     def _check_ip_v6(self, reader, file_name):
         subnets = [
             '::1:ffff:ffff', '::2:0:0', '::2:0:40', '::2:0:50', '::2:0:58'
         ]
 
         for address in subnets:
-            self.assertEqual({'ip': address}, reader.get(address),
+            self.assertEqual({'ip': address}, reader.get(self.ipf(address)),
                              'found expected data record for ' + address +
                              ' in ' + file_name)
 
@@ -386,12 +386,13 @@ def _check_ip_v6(self, reader, file_name):
         }
 
         for key_address, value_address in pairs.items():
-            self.assertEqual({'ip': value_address}, reader.get(key_address),
+            self.assertEqual({'ip': value_address},
+                             reader.get(self.ipf(key_address)),
                              'found expected data record for ' + key_address +
                              ' in ' + file_name)
 
         for ip in ['1.1.1.33', '255.254.253.123', '89fa::']:
-            self.assertIsNone(reader.get(ip))
+            self.assertIsNone(reader.get(self.ipf(ip)))
 
 
 def has_maxminddb_extension():
@@ -408,6 +409,17 @@ class TestExtensionReader(BaseTestReader, unittest.TestCase):
         readerClass = [maxminddb.extension.Reader]
 
 
+@unittest.skipIf(not has_maxminddb_extension()
+                 and not os.environ.get('MM_FORCE_EXT_TESTS'),
+                 'No C extension module found. Skipping tests')
+class TestExtensionReaderWithIPObjects(BaseTestReader, unittest.TestCase):
+    mode = MODE_MMAP_EXT
+    use_ip_objects = True
+
+    if has_maxminddb_extension():
+        readerClass = [maxminddb.extension.Reader]
+
+
 class TestAutoReader(BaseTestReader, unittest.TestCase):
     mode = MODE_AUTO
 
@@ -422,6 +434,14 @@ class TestMMAPReader(BaseTestReader, unittest.TestCase):
     readerClass = [maxminddb.reader.Reader]
 
 
+# We want one pure Python test to use IP objects, it doesn't
+# really matter which one.
+class TestMMAPReaderWithIPObjects(BaseTestReader, unittest.TestCase):
+    mode = MODE_MMAP
+    use_ip_objects = True
+    readerClass = [maxminddb.reader.Reader]
+
+
 class TestFileReader(BaseTestReader, unittest.TestCase):
     mode = MODE_FILE
     readerClass = [maxminddb.reader.Reader]