74 lines
2.6 KiB
Python
74 lines
2.6 KiB
Python
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)
|