Add IP resolver utility that returns whether an IP is under AWS

util/ipresolver/test/test_ipresolver.py

@@ -0,0 +1,50 @@
+import requests
+import pytest
+import json
+
+from httmock import urlmatch, HTTMock
+
+from config import build_requests_session
+from util.ipresolver import IPResolver, ResolvedLocation
+from test.fixtures import *
+
+@pytest.fixture(scope="module")
+def http_client():
+  sess = requests.Session()
+  adapter = requests.adapters.HTTPAdapter(pool_connections=100,
+                                          pool_maxsize=100)
+  sess.mount('http://', adapter)
+  sess.mount('https://', adapter)
+  return sess
+
+@pytest.fixture()
+def aws_ip_range_handler():
+  @urlmatch(netloc=r'ip-ranges.amazonaws.com')
+  def handler(_, request):
+    fake_range_doc = {
+      'syncToken': 123456789,
+      'prefixes': [
+        {
+          'ip_prefix': '10.0.0.0/8',
+          'region': 'GLOBAL',
+          'service': 'AMAZON',
+        } 
+      ],
+    }
+    return {'status_code': 200, 'content': json.dumps(fake_range_doc)}
+
+  return handler
+
+def test_unstarted(app, http_client):
+  ipresolver = IPResolver(app, client=http_client)
+  assert ipresolver.resolve_ip('10.0.0.1') is None
+
+def test_resolved(aws_ip_range_handler, app, http_client):
+  with HTTMock(aws_ip_range_handler):
+    ipresolver = IPResolver(app, client=http_client)
+    assert ipresolver._update_aws_ip_range()
+
+    assert ipresolver.resolve_ip('10.0.0.1') == ResolvedLocation(provider='aws', region=u'GLOBAL', service=None, sync_token=123456789)
+    assert ipresolver.resolve_ip('10.0.0.2') == ResolvedLocation(provider='aws', region=u'GLOBAL', service=None, sync_token=123456789)
+    assert ipresolver.resolve_ip('1.2.3.4') == ResolvedLocation(provider='internet', region=u'NA', service=u'US', sync_token=123456789)
+    assert ipresolver.resolve_ip('127.0.0.1') == ResolvedLocation(provider='internet', region=None, service=None, sync_token=123456789)