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)