quay/util/ipresolver/test/test_ipresolver.py

import pytest

from mock import patch

from util.ipresolver import _UpdateIPRange, IPResolver, ResolvedLocation, CACHE, _FAILED_UPDATE_RETRY_SECS, _UPDATE_INTERVAL
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',
      },
      {
        'ip_prefix': '6.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 not 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, country_iso_code=None)
    assert ipresolver.resolve_ip('10.0.0.2') == ResolvedLocation(provider='aws', region=u'GLOBAL', service=None, sync_token=123456789, country_iso_code=None)
    assert ipresolver.resolve_ip('6.0.0.2') == ResolvedLocation(provider='aws', region=u'GLOBAL', service=None, sync_token=123456789, country_iso_code=u'US')
    assert ipresolver.resolve_ip('1.2.3.4') == ResolvedLocation(provider='internet', region=u'NA', service=u'US', sync_token=123456789, country_iso_code=u'US')
    assert ipresolver.resolve_ip('127.0.0.1') == ResolvedLocation(provider='internet', region=None, service=None, sync_token=123456789, country_iso_code=None)

def test_thread_missing_file():
  class LoopInterruptionForTest(Exception):
    pass

  with patch('util.ipresolver._get_aws_ip_ranges', return_value=None):
    with patch('time.sleep', side_effect=LoopInterruptionForTest) as sleep:
      update_ip_thread = _UpdateIPRange(_UPDATE_INTERVAL)
      update_ip_thread.start()
      update_ip_thread.join(30)
      sleep.assert_called_with(_FAILED_UPDATE_RETRY_SECS)