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 test_aws_ip(): return '10.0.0.1' @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, test_aws_ip, http_client): ipresolver = IPResolver(app, client=http_client) assert ipresolver.resolve_ip(test_aws_ip) is None def test_resolved(aws_ip_range_handler, test_aws_ip, 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(test_aws_ip) == 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)