quay/util/ipresolver/test/test_ipresolver.py

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)