quay/util/ipresolver/test/test_ipresolver.py

import pytest

from mock import patch

from util.ipresolver import IPResolver, ResolvedLocation
from test.fixtures import *

@pytest.fixture()
def test_aws_ip():
  return '10.0.0.1'

@pytest.fixture()
def aws_ip_range_data():
  fake_range_doc = {
    'syncToken': 123456789,
    'prefixes': [
      {
        'ip_prefix': '10.0.0.0/8',
        'region': 'GLOBAL',
        'service': 'AMAZON',
      } 
    ],
  }
  return fake_range_doc

def test_unstarted(app, test_aws_ip):
  ipresolver = IPResolver(app)
  assert ipresolver.resolve_ip(test_aws_ip) is None

def test_resolved(aws_ip_range_data, test_aws_ip, app,):
    ipresolver = IPResolver(app)

    def get_data():
      return aws_ip_range_data

    with patch.object(ipresolver, '_get_aws_ip_ranges', get_data):
      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)