def test_arctic_auth():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True), \
patch('arctic.arctic.mongo_retry', autospec=True), \
patch('arctic.arctic.get_auth', autospec=True) as ga:
ga.return_value = Credential('db', 'admin_user', 'admin_pass')
store = Arctic('cluster')
# do something to trigger lazy arctic init
store.list_libraries()
ga.assert_called_once_with('cluster', 'arctic', 'admin')
store._adminDB.authenticate.assert_called_once_with(
'admin_user', 'admin_pass')
ga.reset_mock()
# Get a 'missing' library
with pytest.raises(LibraryNotFoundException):
with patch('arctic.arctic.ArcticLibraryBinding.get_library_type',
return_value=None,
autospec=True):
ga.return_value = Credential('db', 'user', 'pass')
store._conn['arctic_jblackburn'].name = 'arctic_jblackburn'
store['jblackburn.library']
# Creating the library will have attempted to auth against it
ga.assert_called_once_with('cluster', 'arctic', 'arctic_jblackburn')
store._conn['arctic_jblackburn'].authenticate.assert_called_once_with(
'user', 'pass')
class Appender(object):
def __init__(self, mongo_server, library_name, sem, counter_init, runtime=30):
super(Appender, self).__init__()
self.lib = Arctic(mongo_server)[library_name]
self.sem = sem
self.begin = counter_init
self.last = counter_init
self.timeout = datetime.now() + timedelta(seconds=runtime)
def run(self):
self.sem.acquire()
while datetime.now() < self.timeout:
try:
# Randomy length dataframe to keep appending to
df = DataFrame({'v': [self.last]}, [datetime.now()])
for i in range(random.randint(1, 10)):
df = df.append(DataFrame({'v': [self.last + i]}, [datetime.now()]))
self.last + i
df.index.name = 'index'
self.lib.append('symbol', df)
assert self.last in self.lib.read('symbol').data['v'].tolist()
self.last += 2
except OptimisticLockException:
# Concurrent write, not successful
pass
# time.sleep(self.begin)
def check_written_data_exists(self):
values = self.lib.read('symbol').data['v'].tolist()
assert len(set(values)) == len(values), "Written: %s" % values
i = self.begin
while i < self.last:
assert i in values, "Missing %s in %s" % (i, values)
i += 2
def test_get_library_not_initialized():
self = create_autospec(Arctic, mongo_host=sentinel.host)
self._library_cache = {}
with pytest.raises(LibraryNotFoundException) as e, patch("arctic.arctic.ArcticLibraryBinding", autospec=True) as ML:
ML.return_value.get_library_type.return_value = None
Arctic.get_library(self, sentinel.lib_name)
assert "Library %s was not correctly initialized in %s." % (sentinel.lib_name, self) in str(e)
def test_arctic_auth_custom_app_name():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True), \
patch('arctic.arctic.mongo_retry', autospec=True), \
patch('arctic._cache.Cache._is_not_expired', return_value=True), \
patch('arctic.arctic.get_auth', autospec=True) as ga:
ga.return_value = Credential('db', 'admin_user', 'admin_pass')
store = Arctic('cluster', app_name=sentinel.app_name)
# do something to trigger lazy arctic init
store.list_libraries()
assert ga.call_args_list == [
call('cluster', sentinel.app_name, 'admin')
]
ga.reset_mock()
# Get a 'missing' library
with pytest.raises(LibraryNotFoundException):
with patch('arctic.arctic.ArcticLibraryBinding.get_library_type',
return_value=None,
autospec=True):
ga.return_value = Credential('db', 'user', 'pass')
store._conn['arctic_jblackburn'].name = 'arctic_jblackburn'
store['jblackburn.library']
# Creating the library will have attempted to auth against it
assert ga.call_args_list == [
call('cluster', sentinel.app_name, 'arctic_jblackburn')
]
def test_initialize_library_with_list_coll_names():
self = create_autospec(Arctic)
self._conn = create_autospec(MongoClient)
lib = create_autospec(ArcticLibraryBinding)
lib.database_name = sentinel.db_name
lib.get_quota.return_value = None
self._conn.__getitem__.return_value.list_collection_names.return_value = [
x for x in six.moves.xrange(5001)
]
lib_type = Mock()
with patch.dict('arctic.arctic.LIBRARY_TYPES', {sentinel.lib_type: lib_type}), \
patch('arctic.arctic.ArcticLibraryBinding', return_value=lib, autospec=True) as ML:
Arctic.initialize_library(self,
sentinel.lib_name,
sentinel.lib_type,
thing=sentinel.thing,
check_library_count=False)
assert ML.call_args_list == [call(self, sentinel.lib_name)]
assert ML.return_value.set_library_type.call_args_list == [
call(sentinel.lib_type)
]
assert ML.return_value.set_quota.call_args_list == [
call(10 * 1024 * 1024 * 1024)
]
assert lib_type.initialize_library.call_args_list == [
call(ML.return_value, thing=sentinel.thing)
]
def test_get_library_auth_issue():
self = create_autospec(Arctic, mongo_host=sentinel.host)
self._library_cache = {}
with pytest.raises(LibraryNotFoundException) as e, \
patch('arctic.arctic.ArcticLibraryBinding', autospec=True) as ML:
ML.return_value.get_library_type.side_effect = OperationFailure('database error: not authorized for query on arctic_marketdata.index.ARCTIC')
Arctic.get_library(self, sentinel.lib_name)
assert "Library %s was not correctly initialized in %s." % (sentinel.lib_name, self) in str(e)
def test_get_library_auth_issue():
self = create_autospec(Arctic, mongo_host=sentinel.host)
self._library_cache = {}
with pytest.raises(LibraryNotFoundException) as e, \
patch('arctic.arctic.ArcticLibraryBinding', autospec=True) as ML:
ML.return_value.get_library_type.side_effect = OperationFailure('database error: not authorized for query on arctic_marketdata.index.ARCTIC')
Arctic.get_library(self, sentinel.lib_name)
assert "Library %s was not correctly initialized in %s." % (sentinel.lib_name, self) in str(e)
def test_get_library_not_initialized():
self = create_autospec(Arctic, mongo_host=sentinel.host)
self._library_cache = {}
with pytest.raises(LibraryNotFoundException) as e, \
patch('arctic.arctic.ArcticLibraryBinding', autospec=True) as ML:
ML.return_value.get_library_type.return_value = None
Arctic.get_library(self, sentinel.lib_name)
assert "Library %s was not correctly initialized in %s." % (
sentinel.lib_name, self) in str(e)
def test_arctic_lazy_init():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True) as mc, \
patch('arctic.arctic.mongo_retry', side_effect=lambda x: x, autospec=True), \
patch('arctic.arctic.get_auth', autospec=True) as ga:
store = Arctic('cluster')
assert not mc.called
# do something to trigger lazy arctic init
store.list_libraries()
assert mc.called
def test_arctic_lazy_init():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True) as mc, \
patch('arctic.arctic.mongo_retry', side_effect=lambda x: x, autospec=True), \
patch('arctic.arctic.get_auth', autospec=True) as ga:
store = Arctic('cluster')
assert not mc.called
# do something to trigger lazy arctic init
store.list_libraries()
assert mc.called
def test_get_library_not_registered():
self = create_autospec(Arctic)
self._library_cache = {}
with pytest.raises(LibraryNotFoundException) as e, \
patch('arctic.arctic.ArcticLibraryBinding', autospec=True) as ML:
ML.return_value.get_library_type.return_value = sentinel.lib_type
Arctic.get_library(self, sentinel.lib_name)
assert ("Couldn't load LibraryType '%s' for '%s' (has the class been registered?)" %
(sentinel.lib_type, sentinel.lib_name)
)in str(e)
def __init__(self, connection=None):
if not connection:
connection = 'localhost'
self.__mongo_client = connection
# tick db
Arctic.__init__(self, self.__mongo_client)
# job system
self.__job_db = self.__mongo_client.jobs
self.__job_collection = self.__job_db.queue
def test_get_library_not_registered():
self = create_autospec(Arctic)
self._library_cache = {}
with pytest.raises(LibraryNotFoundException) as e, \
patch('arctic.arctic.ArcticLibraryBinding', autospec=True) as ML:
ML.return_value.get_library_type.return_value = sentinel.lib_type
Arctic.get_library(self, sentinel.lib_name)
assert (
"Couldn't load LibraryType '%s' for '%s' (has the class been registered?)"
% (sentinel.lib_type, sentinel.lib_name)) in str(e)
def __init__(self,
mongo_server,
library_name,
sem,
counter_init,
runtime=30):
super(Appender, self).__init__()
self.lib = Arctic(mongo_server)[library_name]
self.sem = sem
self.begin = counter_init
self.last = counter_init
self.timeout = datetime.now() + timedelta(seconds=runtime)
def test_arctic_connect_hostname():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True) as mc, \
patch('arctic.arctic.mongo_retry', autospec=True) as ar, \
patch('arctic.arctic.get_mongodb_uri', autospec=True) as gmu:
store = Arctic('hostname', socketTimeoutMS=sentinel.socket_timeout,
connectTimeoutMS=sentinel.connect_timeout,
serverSelectionTimeoutMS=sentinel.select_timeout)
# do something to trigger lazy arctic init
store.list_libraries()
mc.assert_called_once_with(host=gmu('hostname'), maxPoolSize=4,
socketTimeoutMS=sentinel.socket_timeout,
connectTimeoutMS=sentinel.connect_timeout,
serverSelectionTimeoutMS=sentinel.select_timeout)
def test_arctic_connect_hostname():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True) as mc, \
patch('arctic.arctic.mongo_retry', autospec=True) as ar, \
patch('arctic.arctic.get_mongodb_uri', autospec=True) as gmu:
store = Arctic('hostname', socketTimeoutMS=sentinel.socket_timeout,
connectTimeoutMS=sentinel.connect_timeout,
serverSelectionTimeoutMS=sentinel.select_timeout)
# do something to trigger lazy arctic init
store.list_libraries()
ar(mc).assert_called_once_with(host=gmu('hostname'), maxPoolSize=4,
socketTimeoutMS=sentinel.socket_timeout,
connectTimeoutMS=sentinel.connect_timeout,
serverSelectionTimeoutMS=sentinel.select_timeout)
def test_initialize_library():
self = create_autospec(Arctic)
self._conn = create_autospec(MongoClient)
lib = create_autospec(ArcticLibraryBinding)
lib.database_name = sentinel.db_name
lib.get_quota.return_value = None
lib_type = Mock()
with patch.dict('arctic.arctic.LIBRARY_TYPES', {sentinel.lib_type: lib_type}), \
patch('arctic.arctic.ArcticLibraryBinding', return_value=lib, autospec=True) as ML:
Arctic.initialize_library(self, sentinel.lib_name, sentinel.lib_type, thing=sentinel.thing)
assert ML.call_args_list == [call(self, sentinel.lib_name)]
assert ML.return_value.set_library_type.call_args_list == [call(sentinel.lib_type)]
assert ML.return_value.set_quota.call_args_list == [call(10 * 1024 * 1024 * 1024)]
assert lib_type.initialize_library.call_args_list == [call(ML.return_value, thing=sentinel.thing)]
def __init__(self, connection = None):
if not connection:
connection = 'localhost'
self.__mongo_client = connection
# tick db
Arctic.__init__(self, self.__mongo_client)
# job system
self.__job_db = self.__mongo_client.jobs
self.__job_collection = self.__job_db.queue
# scheduler
self.__scheduler = BackgroundScheduler()
self.__scheduler.configure(executors = executors, job_defaults= job_defaults, timezone=utc)
self.__scheduler.start()
def test_initialize_library():
self = create_autospec(Arctic)
self._conn = create_autospec(MongoClient)
lib = create_autospec(ArcticLibraryBinding)
lib.database_name = sentinel.db_name
lib.get_quota.return_value = None
lib_type = Mock()
with patch.dict('arctic.arctic.LIBRARY_TYPES', {sentinel.lib_type: lib_type}), \
patch('arctic.arctic.ArcticLibraryBinding', return_value=lib, autospec=True) as ML:
Arctic.initialize_library(self, sentinel.lib_name, sentinel.lib_type, thing=sentinel.thing)
assert ML.call_args_list == [call(self, sentinel.lib_name)]
assert ML.return_value.set_library_type.call_args_list == [call(sentinel.lib_type)]
assert ML.return_value.set_quota.call_args_list == [call(10 * 1024 * 1024 * 1024)]
assert lib_type.initialize_library.call_args_list == [call(ML.return_value, thing=sentinel.thing)]
def test_reset_Arctic(mongo_host, library_name):
arctic = Arctic(mongo_host=mongo_host)
arctic.list_libraries()
arctic.initialize_library(library_name, VERSION_STORE)
arctic[library_name]
c = arctic._conn
arctic.reset()
assert len(c.nodes) == 0
def test_initialize_library_too_many_ns():
self = create_autospec(Arctic)
self._conn = create_autospec(MongoClient)
lib = create_autospec(ArcticLibraryBinding)
lib.database_name = sentinel.db_name
self._conn.__getitem__.return_value.collection_names.return_value = [x for x in six.moves.xrange(5001)]
lib_type = Mock()
with pytest.raises(ArcticException) as e:
with patch.dict('arctic.arctic.LIBRARY_TYPES', {sentinel.lib_type: lib_type}), \
patch('arctic.arctic.ArcticLibraryBinding', return_value=lib, autospec=True) as ML:
Arctic.initialize_library(self, sentinel.lib_name, sentinel.lib_type, thing=sentinel.thing)
assert self._conn.__getitem__.call_args_list == [call(sentinel.db_name),
call(sentinel.db_name)]
assert lib_type.initialize_library.call_count == 0
assert 'Too many namespaces 5001, not creating: sentinel.lib_name' in str(e)
def test_initialize_library():
self = create_autospec(Arctic)
self._conn = create_autospec(MongoClient)
lib = create_autospec(ArcticLibraryBinding)
lib.database_name = sentinel.db_name
lib.get_quota.return_value = None
self._conn.__getitem__.return_value.collection_names.return_value = [x for x in six.moves.xrange(5001)]
lib_type = Mock()
with patch.dict('arctic.arctic.LIBRARY_TYPES', {sentinel.lib_type: lib_type}), \
patch('arctic.arctic.ArcticLibraryBinding', return_value=lib, autospec=True) as ML:
Arctic.initialize_library(self, sentinel.lib_name, sentinel.lib_type, thing=sentinel.thing, check_library_count=False)
assert ML.call_args_list == [call(self, sentinel.lib_name)]
assert ML.return_value.set_library_type.call_args_list == [call(sentinel.lib_type)]
assert ML.return_value.set_quota.call_args_list == [call(10 * 1024 * 1024 * 1024)]
assert lib_type.initialize_library.call_args_list == [call(ML.return_value, thing=sentinel.thing)]
def test_arctic_connect_with_environment_name():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True) as mc, \
patch('arctic.arctic.mongo_retry', autospec=True) as ar, \
patch('arctic.arctic.get_auth', autospec=True), \
patch('arctic.arctic.get_mongodb_uri') as gmfe:
store = Arctic('live', socketTimeoutMS=sentinel.socket_timeout,
connectTimeoutMS=sentinel.connect_timeout,
serverSelectionTimeoutMS=sentinel.select_timeout)
# do something to trigger lazy arctic init
store.list_libraries()
assert gmfe.call_args_list == [call('live')]
assert mc.call_args_list == [call(host=gmfe.return_value, maxPoolSize=4,
socketTimeoutMS=sentinel.socket_timeout,
connectTimeoutMS=sentinel.connect_timeout,
serverSelectionTimeoutMS=sentinel.select_timeout)]
def test_initialize_library_too_many_ns():
self = create_autospec(Arctic)
self._conn = create_autospec(MongoClient)
lib = create_autospec(ArcticLibraryBinding)
lib.database_name = sentinel.db_name
self._conn.__getitem__.return_value.collection_names.return_value = [x for x in xrange(3001)]
lib_type = Mock()
with pytest.raises(ArcticException) as e:
with patch.dict('arctic.arctic.LIBRARY_TYPES', {sentinel.lib_type: lib_type}), \
patch('arctic.arctic.ArcticLibraryBinding', return_value=lib, autospec=True) as ML:
Arctic.initialize_library(self, sentinel.lib_name, sentinel.lib_type, thing=sentinel.thing)
assert self._conn.__getitem__.call_args_list == [call(sentinel.db_name),
call(sentinel.db_name)]
assert lib_type.initialize_library.call_count == 0
assert 'Too many namespaces 3001, not creating: sentinel.lib_name' in str(e)
def test_arctic_connect_with_environment_name():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True) as mc, \
patch('arctic.arctic.mongo_retry', autospec=True) as ar, \
patch('arctic.arctic.get_auth', autospec=True), \
patch('arctic.arctic.get_mongodb_uri') as gmfe:
store = Arctic('live', socketTimeoutMS=sentinel.socket_timeout,
connectTimeoutMS=sentinel.connect_timeout,
serverSelectionTimeoutMS=sentinel.select_timeout)
# do something to trigger lazy arctic init
store.list_libraries()
assert gmfe.call_args_list == [call('live')]
assert ar(mc).call_args_list == [call(host=gmfe.return_value, maxPoolSize=4,
socketTimeoutMS=sentinel.socket_timeout,
connectTimeoutMS=sentinel.connect_timeout,
serverSelectionTimeoutMS=sentinel.select_timeout)]
def __init__(self, mongo_server, library_name, sem, counter_init, runtime=30):
super(Appender, self).__init__()
self.lib = Arctic(mongo_server)[library_name]
self.sem = sem
self.begin = counter_init
self.last = counter_init
self.timeout = datetime.now() + timedelta(seconds=runtime)
def test_arctic_lazy_init_ssl_true():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True) as mc, \
patch('arctic.arctic.mongo_retry', side_effect=lambda x: x, autospec=True), \
patch('arctic.arctic.get_auth', autospec=True) as ga:
store = Arctic('cluster', ssl=True)
assert not mc.called
# do something to trigger lazy arctic init
store.list_libraries()
assert mc.called
assert len(mc.mock_calls) == 1
assert mc.mock_calls[0] == call(connectTimeoutMS=2000,
host='cluster',
maxPoolSize=4,
serverSelectionTimeoutMS=30000,
socketTimeoutMS=600000,
ssl=True)
def test_arctic_set_get_state():
sentinel.mongo_host = Mock(nodes={("host", "port")})
store = Arctic(sentinel.mongo_host, allow_secondary="allow_secondary")
buff = pickle.dumps(store)
mnew = pickle.loads(buff)
assert mnew.mongo_host == "host:port"
assert mnew._allow_secondary == "allow_secondary"
def test_mongo_host_get_set():
sentinel.mongo_host = Mock(nodes={("host", "port")})
with patch('arctic._cache.Cache.__init__',
autospec=True,
return_value=None):
arctic = Arctic(sentinel.mongo_host)
assert arctic.mongo_host == "host:port"
def test_arctic_repr():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True):
with patch('arctic.arctic.mongo_retry', autospec=True):
with patch('arctic.arctic.get_auth', autospec=True) as ga:
ga.return_value = Credential('db', 'admin_user', 'admin_pass')
store = Arctic('cluster')
assert str(store) == repr(store)
def test_arctic_lazy_init_ssl_true():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True) as mc, \
patch('arctic.arctic.mongo_retry', side_effect=lambda x: x, autospec=True), \
patch('arctic.arctic.get_auth', autospec=True) as ga:
store = Arctic('cluster', ssl=True)
assert not mc.called
# do something to trigger lazy arctic init
store.list_libraries()
assert mc.called
assert len(mc.mock_calls) == 1
assert mc.mock_calls[0] == call(connectTimeoutMS=2000,
host='cluster',
maxPoolSize=4,
serverSelectionTimeoutMS=30000,
socketTimeoutMS=600000,
ssl=True)
def test_re_authenticate_on_arctic_reset(mongo_host, library_name):
from collections import namedtuple
Cred = namedtuple('Cred', 'user, password')
with patch('arctic.arctic.authenticate') as auth_mock, \
patch('arctic.arctic.get_auth') as get_auth_mock:
auth_mock.return_value = True
get_auth_mock.return_value = Cred(user='******', password='******')
arctic = Arctic(mongo_host=mongo_host)
arctic.initialize_library(library_name, VERSION_STORE)
vstore = arctic[library_name]
vstore.list_symbols()
auth_mock.reset_mock()
arctic.reset()
assert auth_mock.call_count > 0
auth_mock.reset_mock()
vstore.list_symbols()
assert auth_mock.call_count == 0
def test_reset_Arctic(mongo_host, library_name):
arctic = Arctic(mongo_host=mongo_host)
arctic.list_libraries()
arctic.initialize_library(library_name, VERSION_STORE)
c = arctic._conn
assert arctic[library_name]._arctic_lib._curr_conn is c
arctic.reset()
assert c is not arctic._conn
assert len(c.nodes) == 0
assert arctic[library_name]._arctic_lib._curr_conn is arctic._conn
def test_reset():
c = MagicMock()
with patch('pymongo.MongoClient', return_value=c, autospec=True) as mc:
store = Arctic('hostname')
# do something to trigger lazy arctic init
store.list_libraries()
store.reset()
# Doesn't matter how many times we call it:
store.reset()
c.close.assert_called_once()
def test_init_library_quota(mongo_host):
# Create the user agains the current mongo database
with patch('arctic.scripts.arctic_init_library.do_db_auth', return_value=True), \
patch('pymongo.database.Database.authenticate', return_value=True):
run_as_main(mil.main, '--host', mongo_host, '--library', 'arctic_user.library', '--quota', '100')
# Should be able to write something to the library now
store = Arctic(mongo_host)
assert store['user.library']._arctic_lib.get_library_metadata('QUOTA') == 100 * 1024 * 1024 * 1024
class Appender(object):
def __init__(self,
mongo_server,
library_name,
sem,
counter_init,
runtime=30):
super(Appender, self).__init__()
self.lib = Arctic(mongo_server)[library_name]
self.sem = sem
self.begin = counter_init
self.last = counter_init
self.timeout = datetime.now() + timedelta(seconds=runtime)
def run(self):
self.sem.acquire()
while datetime.now() < self.timeout:
try:
# Randomy length dataframe to keep appending to
df = DataFrame({'v': [self.last]}, [datetime.now()])
for i in range(random.randint(1, 10)):
df = df.append(
DataFrame({'v': [self.last + i]}, [datetime.now()]))
self.last + i
df.index.name = 'index'
self.lib.append('symbol', df)
assert self.last in self.lib.read('symbol').data['v'].tolist()
self.last += 2
except OptimisticLockException:
# Concurrent write, not successful
pass
# time.sleep(self.begin)
def check_written_data_exists(self):
values = self.lib.read('symbol').data['v'].tolist()
assert len(set(values)) == len(values), "Written: %s" % values
i = self.begin
while i < self.last:
assert i in values, "Missing %s in %s" % (i, values)
i += 2
def test_arctic_auth_custom_app_name():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True), \
patch('arctic.arctic.mongo_retry', autospec=True), \
patch('arctic.arctic.get_auth', autospec=True) as ga:
ga.return_value = Credential('db', 'admin_user', 'admin_pass')
store = Arctic('cluster', app_name=sentinel.app_name)
# do something to trigger lazy arctic init
store.list_libraries()
assert ga.call_args_list == [call('cluster', sentinel.app_name, 'admin')]
ga.reset_mock()
# Get a 'missing' library
with pytest.raises(LibraryNotFoundException):
with patch('arctic.arctic.ArcticLibraryBinding.get_library_type', return_value=None, autospec=True):
ga.return_value = Credential('db', 'user', 'pass')
store._conn['arctic_jblackburn'].name = 'arctic_jblackburn'
store['jblackburn.library']
# Creating the library will have attempted to auth against it
assert ga.call_args_list == [call('cluster', sentinel.app_name, 'arctic_jblackburn')]
def test_arctic_auth_custom_app_name():
with patch("pymongo.MongoClient", return_value=MagicMock(), autospec=True), patch(
"arctic.arctic.mongo_retry", autospec=True
), patch("arctic.arctic.get_auth", autospec=True) as ga:
ga.return_value = Credential("db", "admin_user", "admin_pass")
store = Arctic("cluster", app_name=sentinel.app_name)
# do something to trigger lazy arctic init
store.list_libraries()
assert ga.call_args_list == [call("cluster", sentinel.app_name, "admin")]
ga.reset_mock()
# Get a 'missing' library
with pytest.raises(LibraryNotFoundException):
with patch("arctic.arctic.ArcticLibraryBinding.get_library_type", return_value=None, autospec=True):
ga.return_value = Credential("db", "user", "pass")
store._conn["arctic_jblackburn"].name = "arctic_jblackburn"
store["jblackburn.library"]
# Creating the library will have attempted to auth against it
assert ga.call_args_list == [call("cluster", sentinel.app_name, "arctic_jblackburn")]
def test_get_library():
self = create_autospec(Arctic)
self._library_cache = {}
library_type = Mock()
register_library_type(sentinel.lib_type, library_type)
with patch('arctic.arctic.ArcticLibraryBinding', autospec=True) as ML:
ML.return_value.get_library_type.return_value = sentinel.lib_type
library = Arctic.get_library(self, sentinel.lib_name)
del LIBRARY_TYPES[sentinel.lib_type]
assert ML.call_args_list == [call(self, sentinel.lib_name)]
assert library_type.call_args_list == [call(ML.return_value)]
assert library == library_type.return_value
def test_reset():
c = MagicMock()
with patch('pymongo.MongoClient', return_value=c, autospec=True) as mc:
store = Arctic('hostname')
# do something to trigger lazy arctic init
store.list_libraries()
store.reset()
# Doesn't matter how many times we call it:
store.reset()
c.close.assert_called_once()
def test_arctic_set_get_state():
sentinel.mongo_host = Mock(nodes={("host", "port")})
store = Arctic(sentinel.mongo_host, allow_secondary="allow_secondary", app_name="app_name",
socketTimeoutMS=1234, connectTimeoutMS=2345, serverSelectionTimeoutMS=3456)
buff = pickle.dumps(store)
mnew = pickle.loads(buff)
assert mnew.mongo_host == "host:port"
assert mnew._allow_secondary == "allow_secondary"
assert mnew._application_name == "app_name"
assert mnew._socket_timeout == 1234
assert mnew._connect_timeout == 2345
assert mnew._server_selection_timeout == 3456
def test_get_library():
self = create_autospec(Arctic)
self._library_cache = {}
library_type = Mock()
register_library_type(sentinel.lib_type, library_type)
with patch('arctic.arctic.ArcticLibraryBinding', autospec=True) as ML:
ML.return_value.get_library_type.return_value = sentinel.lib_type
library = Arctic.get_library(self, sentinel.lib_name)
del LIBRARY_TYPES[sentinel.lib_type]
assert ML.call_args_list == [call(self, sentinel.lib_name)]
assert library_type.call_args_list == [call(ML.return_value)]
assert library == library_type.return_value
def test_reset_Arctic(mongo_host, library_name):
arctic = Arctic(mongo_host=mongo_host)
arctic.list_libraries()
arctic.initialize_library(library_name, VERSION_STORE)
c = arctic._conn
assert arctic[library_name]._arctic_lib._curr_conn is c
arctic.reset()
assert c is not arctic._conn
assert len(c.nodes) == 0
assert arctic[library_name]._arctic_lib._curr_conn is arctic._conn
def test_arctic_connect_with_environment_name():
with patch("pymongo.MongoClient", return_value=MagicMock(), autospec=True) as mc, patch(
"arctic.arctic.mongo_retry", autospec=True
) as ar, patch("arctic.arctic.get_auth", autospec=True), patch("arctic.arctic.get_mongodb_uri") as gmfe:
store = Arctic(
"live",
socketTimeoutMS=sentinel.socket_timeout,
connectTimeoutMS=sentinel.connect_timeout,
serverSelectionTimeoutMS=sentinel.select_timeout,
)
# do something to trigger lazy arctic init
store.list_libraries()
assert gmfe.call_args_list == [call("live")]
assert ar(mc).call_args_list == [
call(
host=gmfe.return_value,
maxPoolSize=4,
socketTimeoutMS=sentinel.socket_timeout,
connectTimeoutMS=sentinel.connect_timeout,
serverSelectionTimeoutMS=sentinel.select_timeout,
)
]
def test_arctic_set_get_state():
sentinel.mongo_host = Mock(nodes={("host", "port")})
with patch('arctic._cache.Cache.__init__', autospec=True, return_value=None):
store = Arctic(sentinel.mongo_host, allow_secondary="allow_secondary", app_name="app_name",
socketTimeoutMS=1234, connectTimeoutMS=2345, serverSelectionTimeoutMS=3456)
buff = pickle.dumps(store)
mnew = pickle.loads(buff)
assert mnew.mongo_host == "host:port"
assert mnew._allow_secondary == "allow_secondary"
assert mnew._application_name == "app_name"
assert mnew._socket_timeout == 1234
assert mnew._connect_timeout == 2345
assert mnew._server_selection_timeout == 3456
def test_connection_passed_warning_raised():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True), \
patch('arctic.arctic.mongo_retry', side_effect=lambda x: x, autospec=True), \
patch('arctic.arctic.get_auth', autospec=True), \
patch('arctic.arctic.logger') as lg:
magic_mock = MagicMock(nodes={("host", "port")})
store = Arctic(magic_mock, ssl=True)
# Increment _pid to simulate forking the process
store._pid += 1
_ = store._conn
assert lg.mock_calls[0] == call.warn(
'Forking process. Arctic was passed a pymongo connection during init, '
'the new pymongo connection may have different parameters.')
def test_arctic_auth():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True), \
patch('arctic.arctic.mongo_retry', autospec=True), \
patch('arctic.arctic.get_auth', autospec=True) as ga:
ga.return_value = Credential('db', 'admin_user', 'admin_pass')
store = Arctic('cluster')
# do something to trigger lazy arctic init
store.list_libraries()
ga.assert_called_once_with('cluster', 'arctic', 'admin')
store._adminDB.authenticate.assert_called_once_with('admin_user', 'admin_pass')
ga.reset_mock()
# Get a 'missing' library
with pytest.raises(LibraryNotFoundException):
with patch('arctic.arctic.ArcticLibraryBinding.get_library_type', return_value=None, autospec=True):
ga.return_value = Credential('db', 'user', 'pass')
store._conn['arctic_jblackburn'].name = 'arctic_jblackburn'
store._conn['arctic_jblackburn'].command.return_value = mock_conn_info
store['jblackburn.library']
# Creating the library will have attempted to auth against it
ga.assert_called_once_with('cluster', 'arctic', 'arctic_jblackburn')
store._conn['arctic_jblackburn'].authenticate.assert_called_once_with('user', 'pass')
def test_arctic_auth_admin_reauth():
with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True), \
patch('arctic.arctic.mongo_retry', autospec=True), \
patch('arctic.arctic.get_auth', autospec=True) as ga:
ga.return_value = Credential('db', 'admin_user', 'admin_pass')
store = Arctic('cluster')
# do something to trigger lazy arctic init
store.list_libraries()
assert ga.call_args_list == [call('cluster', 'arctic', 'admin')]
ga.reset_mock()
# Get a 'missing' library
with pytest.raises(LibraryNotFoundException):
with patch('arctic.arctic.ArcticLibraryBinding.get_library_type', return_value=None, autospec=True), \
patch('arctic.arctic.logger') as logger:
ga.return_value = Credential('db', 'user', 'pass')
store._conn['arctic_jblackburn'].name = 'arctic_jblackburn'
store._conn['arctic_jblackburn'].command.return_value = mock_conn_info_empty
store['jblackburn.library']
assert store._conn['arctic_jblackburn'].command.call_args_list == [call({'connectionStatus': 1})]
assert ga.call_args_list == [call('cluster', 'arctic', 'arctic_jblackburn'),
call('cluster', 'arctic', store._adminDB.name)]
def test__conn_auth_issue():
auth_timeout = [0]
a = Arctic("host:12345")
sentinel.creds = Mock()
def flaky_auth(*args, **kwargs):
if not auth_timeout[0]:
auth_timeout[0] = 1
raise AutoReconnect()
with patch('arctic.arctic.authenticate', flaky_auth), \
patch('arctic.arctic.get_auth', return_value=sentinel.creds), \
patch('arctic.decorators._handle_error') as he:
a._conn
assert he.call_count == 1
assert auth_timeout[0]
def test_re_authenticate_on_arctic_reset(mongo_host, library_name):
from collections import namedtuple
Cred = namedtuple('Cred', 'user, password')
with patch('arctic.arctic.authenticate') as auth_mock, \
patch('arctic.arctic.get_auth') as get_auth_mock:
auth_mock.return_value = True
get_auth_mock.return_value = Cred(user='******', password='******')
arctic = Arctic(mongo_host=mongo_host)
arctic.initialize_library(library_name, VERSION_STORE)
vstore = arctic[library_name]
vstore.list_symbols()
auth_mock.reset_mock()
arctic.reset()
assert auth_mock.call_count > 0
auth_mock.reset_mock()
vstore.list_symbols()
assert auth_mock.call_count == 0
def test_connect_to_Arctic_connection(mongodb, mongo_host):
arctic = Arctic(mongodb)
assert arctic.list_libraries() == []
assert arctic.mongo_host == mongo_host
def test_connect_to_Arctic_connection(mongodb, mongo_host):
arctic = Arctic(mongodb)
assert arctic.list_libraries() == []
assert arctic.mongo_host == mongo_host
def test_connect_to_Arctic_string(mongo_host):
arctic = Arctic(mongo_host=mongo_host)
assert arctic.list_libraries() == []
assert arctic.mongo_host == mongo_host
def test_library_exists():
self = create_autospec(Arctic)
self.get_library.return_value = 'not an exception'
assert Arctic.library_exists(self, 'mylib')
def test_library_doesnt_exist():
self = create_autospec(Arctic)
self.get_library.side_effect = LibraryNotFoundException('not found')
assert not Arctic.library_exists(self, 'mylib')