import os
import sys
import mock
import time
import random
import shutil
import contextlib
import tempfile
import binascii
import platform
import select
import datetime
from subprocess import Popen, PIPE
from dateutil.tz import tzlocal
if sys.version_info[:2] == (2, 6):
import unittest2 as unittest
else:
import unittest
import botocore.loaders
import botocore.session
from botocore import utils
from botocore import credentials
_LOADER = botocore.loaders.Loader()
def skip_unless_has_memory_collection(cls):
if platform.system() not in ['Darwin', 'Linux']:
return unittest.skip('Memory tests only supported on mac/linux.')(cls)
return cls
def random_chars(num_chars):
return binascii.hexlify(os.urandom(int(num_chars / 2))).decode('ascii')
def create_session(**kwargs):
session = botocore.session.Session(**kwargs)
session.register_component('data_loader', _LOADER)
session.set_config_variable('credentials_file', 'noexist/foo/botocore')
return session
@contextlib.contextmanager
def temporary_file(mode):
temporary_directory = tempfile.mkdtemp()
basename = 'tmpfile-%s-%s' % (int(time.time()), random.randint(1, 1000))
full_filename = os.path.join(temporary_directory, basename)
open(full_filename, 'w').close()
try:
with open(full_filename, mode) as f:
yield f
finally:
shutil.rmtree(temporary_directory)
class BaseEnvVar(unittest.TestCase):
def setUp(self):
self.environ = {}
self.environ_patch = mock.patch('os.environ', self.environ)
self.environ_patch.start()
def tearDown(self):
self.environ_patch.stop()
class BaseSessionTest(BaseEnvVar):
def setUp(self, **environ):
super(BaseSessionTest, self).setUp()
self.environ['AWS_ACCESS_KEY_ID'] = 'access_key'
self.environ['AWS_SECRET_ACCESS_KEY'] = 'secret_key'
self.environ['AWS_CONFIG_FILE'] = 'no-exist-foo'
self.environ.update(environ)
self.session = create_session()
self.session.config_filename = 'no-exist-foo'
@skip_unless_has_memory_collection
class BaseClientDriverTest(unittest.TestCase):
INJECT_DUMMY_CREDS = False
def setUp(self):
self.driver = ClientDriver()
env = None
if self.INJECT_DUMMY_CREDS:
env = {'AWS_ACCESS_KEY_ID': 'foo',
'AWS_SECRET_ACCESS_KEY': 'bar'}
self.driver.start(env=env)
def cmd(self, *args):
self.driver.cmd(*args)
def send_cmd(self, *args):
self.driver.send_cmd(*args)
def record_memory(self):
self.driver.record_memory()
@property
def memory_samples(self):
return self.driver.memory_samples
def tearDown(self):
self.driver.stop()
class ClientDriver(object):
CLIENT_SERVER = os.path.join(
os.path.dirname(os.path.abspath(__file__)),
'cmd-runner'
)
def __init__(self):
self._popen = None
self.memory_samples = []
def _get_memory_with_ps(self, pid):
command_list = ['ps', '-p', str(pid), '-o', 'rss']
p = Popen(command_list, stdout=PIPE)
stdout = p.communicate()[0]
if not p.returncode == 0:
raise RuntimeError("Could not retrieve memory")
else:
return int(stdout.splitlines()[1].split()[0]) * 1024
def record_memory(self):
mem = self._get_memory_with_ps(self._popen.pid)
self.memory_samples.append(mem)
def start(self, env=None):
self._popen = Popen([sys.executable, self.CLIENT_SERVER],
stdout=PIPE, stdin=PIPE, env=env)
def stop(self):
self.cmd('exit')
self._popen.wait()
def send_cmd(self, *cmd):
cmd_str = ' '.join(cmd) + '\n'
cmd_bytes = cmd_str.encode('utf-8')
self._popen.stdin.write(cmd_bytes)
self._popen.stdin.flush()
def is_cmd_finished(self):
rlist = [self._popen.stdout.fileno()]
result = select.select(rlist, [], [], 0.01)
if result[0]:
return True
return False
def cmd(self, *cmd):
self.send_cmd(*cmd)
result = self._popen.stdout.readline().strip()
if result != b'OK':
raise RuntimeError(
"Error from command '%s': %s" % (cmd, result))
class IntegerRefresher(credentials.RefreshableCredentials):
_advisory_refresh_timeout = 2
_mandatory_refresh_timeout = 1
_credentials_expire = 3
def __init__(self, creds_last_for=_credentials_expire,
advisory_refresh=_advisory_refresh_timeout,
mandatory_refresh=_mandatory_refresh_timeout,
refresh_function=None):
expires_in = (
self._current_datetime() +
datetime.timedelta(seconds=creds_last_for))
if refresh_function is None:
refresh_function = self._do_refresh
super(IntegerRefresher, self).__init__(
'0', '0', '0', expires_in,
refresh_function, 'INTREFRESH')
self.creds_last_for = creds_last_for
self.refresh_counter = 0
self._advisory_refresh_timeout = advisory_refresh
self._mandatory_refresh_timeout = mandatory_refresh
def _do_refresh(self):
self.refresh_counter += 1
current = int(self._access_key)
next_id = str(current + 1)
return {
'access_key': next_id,
'secret_key': next_id,
'token': next_id,
'expiry_time': self._seconds_later(self.creds_last_for),
}
def _seconds_later(self, num_seconds):
num_seconds += 1
t = self._current_datetime() + datetime.timedelta(seconds=num_seconds)
return self._to_timestamp(t)
def _to_timestamp(self, datetime_obj):
obj = utils.parse_to_aware_datetime(datetime_obj)
return obj.strftime('%Y-%m-%dT%H:%M:%SZ')
def _current_timestamp(self):
return self._to_timestamp(self._current_datetime())
def _current_datetime(self):
return datetime.datetime.now(tzlocal())