quay/util/ipresolver/test/test_ipresolver.py

import pytest

from mock import patch

from util.ipresolver import IPResolver, ResolvedLocation, CACHE
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

@pytest.fixture()
def test_ip_range_cache(aws_ip_range_data):
  sync_token = aws_ip_range_data['syncToken']
  all_amazon, regions = IPResolver._parse_amazon_ranges(aws_ip_range_data)
  fake_cache = {
    'sync_token': sync_token,
    'all_amazon': all_amazon,
    'regions': regions,
  }
  return fake_cache

@pytest.fixture()
def unstarted_cache():
  fake_unstarted_cache = {}
  return fake_unstarted_cache

def test_unstarted(app, test_aws_ip, unstarted_cache):
  with patch('util.ipresolver._UpdateIPRange'):
    ipresolver = IPResolver(app)

  with patch.dict('util.ipresolver.CACHE', unstarted_cache):
    assert ipresolver.resolve_ip(test_aws_ip) is None

def test_resolved(aws_ip_range_data, test_ip_range_cache, test_aws_ip, app):
  with patch('util.ipresolver._UpdateIPRange'):
    ipresolver = IPResolver(app)

  with patch.dict('util.ipresolver.CACHE', test_ip_range_cache):
    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)