add geo support for ns1

octodns/provider/ns1.py

@@ -6,11 +6,13 @@ from __future__ import absolute_import, division, print_function, \
     unicode_literals
 
 from logging import getLogger
-from nsone import NSONE
+from itertools import chain
+from nsone import NSONE, Config
 from nsone.rest.errors import RateLimitException, ResourceException
+from incf.countryutils import transformations
 from time import sleep
 
-from ..record import Record
+from ..record import _GeoMixin, Record
 from .base import BaseProvider
 
 
@@ -35,11 +37,38 @@ class Ns1Provider(BaseProvider):
         self._client = NSONE(apiKey=api_key)
 
     def _data_for_A(self, _type, record):
-        return {
+        # record meta (which would include geo information is only
+        # returned when getting a record's detail, not from zone detail
+        geo = {}
+        data = {
             'ttl': record['ttl'],
             'type': _type,
-            'values': record['short_answers'],
         }
+        values, codes = [], []
+        if 'answers' not in record:
+            values = record['short_answers']
+        for answer in record.get('answers', []):
+            meta = answer.get('meta', {})
+            if meta:
+                country = meta.get('country', [])
+                us_state = meta.get('us_state', [])
+                ca_province = meta.get('ca_province', [])
+                for cntry in country:
+                    cn = transformations.cc_to_cn(cntry)
+                    con = transformations.cn_to_ctca2(cn)
+                    geo['{}-{}'.format(con, cntry)] = answer['answer']
+                for state in us_state:
+                    geo['NA-US-{}'.format(state)] = answer['answer']
+                for province in ca_province:
+                    geo['NA-CA-{}'.format(state)] = answer['answer']
+                for code in meta.get('iso_region_code', []):
+                    geo[code] = answer['answer']
+            else:
+                values.extend(answer['answer'])
+                codes.append([])
+        data['values'] = values
+        data['geo'] = geo
+        return data
 
     _data_for_AAAA = _data_for_A
 
@@ -146,20 +175,25 @@ class Ns1Provider(BaseProvider):
         try:
             nsone_zone = self._client.loadZone(zone.name[:-1])
             records = nsone_zone.data['records']
+            geo_records = nsone_zone.search(has_geo=True)
         except ResourceException as e:
             if e.message != self.ZONE_NOT_FOUND_MESSAGE:
                 raise
             records = []
+            geo_records = []
 
         before = len(zone.records)
-        for record in records:
+        # geo information isn't returned from the main endpoint, so we need
+        # to query for all records with geo information
+        zone_hash = {}
+        for record in chain(records, geo_records):
             _type = record['type']
             data_for = getattr(self, '_data_for_{}'.format(_type))
             name = zone.hostname_from_fqdn(record['domain'])
             record = Record.new(zone, name, data_for(_type, record),
                                 source=self, lenient=lenient)
-            zone.add_record(record)
-
+            zone_hash[(_type, name)] = record
+        [zone.add_record(r) for r in zone_hash.values()]
         self.log.info('populate:   found %s records',
                       len(zone.records) - before)
 
@@ -168,15 +202,18 @@ class Ns1Provider(BaseProvider):
         if hasattr(record, 'geo'):
             # purposefully set non-geo answers to have an empty meta,
             # so that we know we did this on purpose if/when troubleshooting
-            params['answers'] = [{"answer": x, "meta": {}}
+            params['answers'] = [{"answer": [x], "meta": {}} \
                                  for x in record.values]
             for iso_region, target in record.geo.items():
+                key = 'iso_region_code'
+                value = iso_region
                 params['answers'].append(
                     {
                         'answer': target.values,
-                        'meta': {'iso_region_code': [iso_region]},
+                        'meta': {key: [value]},
                     },
                 )
+        self.log.info("params for A: %s", params)
         return params
 
     _params_for_AAAA = _params_for_A

tests/test_octodns_provider_ns1.py

@@ -38,6 +38,13 @@ class TestNs1Provider(TestCase):
         'values': ['1.2.3.4', '1.2.3.5'],
         'meta': {},
     }))
+    expected.add(Record.new(zone, 'geo', {
+        'ttl': 34,
+        'type': 'A',
+        'values': ['101.102.103.104', '101.102.103.105'],
+        'geo': {'NA-US-NY': ['201.202.203.204']},
+        'meta': {},
+    }))
     expected.add(Record.new(zone, 'cname', {
         'ttl': 34,
         'type': 'CNAME',
@@ -118,6 +125,11 @@ class TestNs1Provider(TestCase):
         'ttl': 33,
         'short_answers': ['1.2.3.4', '1.2.3.5'],
         'domain': 'foo.unit.tests.',
+    }, {
+        'type': 'A',
+        'ttl': 34,
+        'short_answers': ['101.102.103.104', '101.102.103.105'],
+        'domain': 'geo.unit.tests.',
     }, {
         'type': 'CNAME',
         'ttl': 34,
@@ -192,6 +204,9 @@ class TestNs1Provider(TestCase):
         load_mock.reset_mock()
         nsone_zone = DummyZone([])
         load_mock.side_effect = [nsone_zone]
+        zone_search = Mock()
+        zone_search.return_value = []
+        nsone_zone.search = zone_search
         zone = Zone('unit.tests.', [])
         provider.populate(zone)
         self.assertEquals(set(), zone.records)
@@ -201,6 +216,9 @@ class TestNs1Provider(TestCase):
         load_mock.reset_mock()
         nsone_zone = DummyZone(self.nsone_records)
         load_mock.side_effect = [nsone_zone]
+        zone_search = Mock()
+        zone_search.return_value = []
+        nsone_zone.search = zone_search
         zone = Zone('unit.tests.', [])
         provider.populate(zone)
         self.assertEquals(self.expected, zone.records)
@@ -266,11 +284,14 @@ class TestNs1Provider(TestCase):
         }])
         nsone_zone.data['records'][0]['short_answers'][0] = '2.2.2.2'
         nsone_zone.loadRecord = Mock()
+        zone_search = Mock()
+        zone_search.return_value = []
+        nsone_zone.search = zone_search
         load_mock.side_effect = [nsone_zone, nsone_zone]
         plan = provider.plan(desired)
-        self.assertEquals(2, len(plan.changes))
+        self.assertEquals(3, len(plan.changes))
         self.assertIsInstance(plan.changes[0], Update)
-        self.assertIsInstance(plan.changes[1], Delete)
+        self.assertIsInstance(plan.changes[2], Delete)
         # ugh, we need a mock record that can be returned from loadRecord for
         # the update and delete targets, we can add our side effects to that to
         # trigger rate limit handling
@@ -278,23 +299,30 @@ class TestNs1Provider(TestCase):
         mock_record.update.side_effect = [
             RateLimitException('one', period=0),
             None,
+            None,
         ]
         mock_record.delete.side_effect = [
             RateLimitException('two', period=0),
             None,
+            None,
         ]
-        nsone_zone.loadRecord.side_effect = [mock_record, mock_record]
+        nsone_zone.loadRecord.side_effect = [mock_record, mock_record, mock_record]
         got_n = provider.apply(plan)
-        self.assertEquals(2, got_n)
+        self.assertEquals(3, got_n)
         nsone_zone.loadRecord.assert_has_calls([
             call('unit.tests', u'A'),
+            call('geo', u'A'),
             call('delete-me', u'A'),
         ])
         mock_record.assert_has_calls([
-            call.update(answers=[{'answer': u'1.2.3.4', 'meta': {}}], ttl=32),
+            call.update(answers=[{'answer': [u'1.2.3.4'], 'meta': {}}], ttl=32),
+            call.update(answers=[{u'answer': [u'1.2.3.4'], u'meta': {}}], ttl=32),
+            call.update(answers=[{u'answer': [u'101.102.103.104'], u'meta': {}}, {u'answer': [u'101.102.103.105'], u'meta': {}}, {u'answer': [u'201.202.203.204'], u'meta': {u'iso_region_code': [u'NA-US-NY']}}], ttl=34),
+            call.delete(),
             call.delete()
         ])
 
+
     def test_escaping(self):
         provider = Ns1Provider('test', 'api-key')