def test_ArticTransaction_no_audit():
vs = Mock(spec=VersionStore)
ts1 = pd.DataFrame(index=[1, 2], data={'a': [1.0, 2.0]})
vs.read.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=1,
metadata=None,
data=ts1)
vs.write.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=2,
metadata=None,
data=None)
vs.list_versions.return_value = [{'version': 2}, {'version': 1}]
with ArcticTransaction(vs,
sentinel.symbol,
sentinel.user,
sentinel.log,
audit=False) as cwb:
cwb.write(sentinel.symbol,
pd.DataFrame(index=[3, 4], data={'a': [1.0, 2.0]}),
metadata=sentinel.meta)
assert vs.write.call_count == 1
assert vs._write_audit.call_count == 0
def test_ArcticTransaction_writes_if_base_data_corrupted():
vs = Mock(spec=VersionStore)
ts1 = pd.DataFrame(index=[1, 2], data={'a': [1.0, 2.0]})
vs.read.side_effect = OperationFailure('some failure')
vs.write.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=2,
metadata=None,
data=None)
vs.read_metadata.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=1,
metadata=None,
data=None)
vs.list_versions.return_value = [{'version': 2}, {'version': 1}]
with ArcticTransaction(vs, sentinel.symbol, sentinel.user,
sentinel.log) as cwb:
cwb.write(sentinel.symbol, ts1, metadata={1: 2})
vs.write.assert_called_once_with(sentinel.symbol,
ANY,
prune_previous_version=True,
metadata={1: 2})
assert vs.list_versions.call_args_list == [call(sentinel.symbol)]
def test_restore_version():
vs = _create_mock_versionstore()
LASTEST_VERSION = dict(TPL_VERSION, version=TPL_VERSION['version']+1, metadata={'something': 'different'})
last_item = VersionedItem(symbol=TEST_SYMBOL, library=vs._arctic_lib.get_name(),
host=vs._arctic_lib.arctic.mongo_host,
version=LASTEST_VERSION, metadata=LASTEST_VERSION['metadata'], data="hello world")
new_version = dict(LASTEST_VERSION, version=LASTEST_VERSION['version'] + 1)
new_item = VersionedItem(symbol=TEST_SYMBOL, library=vs._arctic_lib.get_name(),
host=vs._arctic_lib.arctic.mongo_host,
version=new_version, metadata=new_version['metadata'], data=last_item.data)
vs.write.return_value = new_item
vs.read.return_value = last_item
vs._read_metadata.side_effect = [TPL_VERSION, LASTEST_VERSION]
with patch('arctic.store.version_store.bson.ObjectId') as mock_objId, \
patch('arctic.store.version_store.mongo_retry') as mock_retry:
mock_objId.return_value = MOCK_OBJID
mock_retry.side_effect = lambda f: f
ret_item = VersionStore.restore_version(vs, symbol=TEST_SYMBOL, as_of=LASTEST_VERSION['version'], prune_previous_version=True)
assert ret_item == new_item
assert vs._read_metadata.call_args_list == [call(TEST_SYMBOL, as_of=LASTEST_VERSION['version'])]
assert vs._version_nums.find_one.call_args_list == [call({'symbol': TEST_SYMBOL})]
assert vs.read.call_args_list == [call(TEST_SYMBOL, as_of=LASTEST_VERSION['version'])]
assert vs.write.call_args_list == [call(TEST_SYMBOL, data=last_item.data, metadata=last_item.metadata, prune_previous_version=True)]
def test_ArcticTransaction_simple():
vs = Mock(spec=VersionStore)
ts1 = pd.DataFrame(index=[1, 2], data={'a': [1.0, 2.0]})
vs.read.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=1,
metadata=None,
data=ts1)
vs.write.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=2,
metadata=None,
data=None)
vs.list_versions.return_value = [{'version': 2}, {'version': 1}]
with ArcticTransaction(vs, sentinel.symbol, sentinel.user,
sentinel.log) as cwb:
cwb.write(sentinel.symbol,
pd.DataFrame(index=[3, 4], data={'a': [1.0, 2.0]}),
metadata=sentinel.meta)
assert not vs._delete_version.called
assert vs.write.call_args_list == [
call(sentinel.symbol,
ANY,
prune_previous_version=True,
metadata=sentinel.meta)
]
assert vs.list_versions.call_args_list == [call(sentinel.symbol)]
assert vs._write_audit.call_args_list == [
call(sentinel.user, sentinel.log, ANY)
]
def test_ConcurrentWriteBlock_simple():
vs = create_autospec(VersionStore, _collection=Mock())
ts1 = pd.DataFrame(index=[1, 2], data={'a': [1.0, 2.0]})
vs.read.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=1,
metadata=None,
data=ts1)
vs.write.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=2,
metadata=None,
data=None)
vs.list_versions.return_value = [{'version': 2}, {'version': 1}]
with ArcticTransaction(vs, sentinel.symbol, sentinel.user,
sentinel.log) as cwb:
cwb.write(sentinel.symbol,
pd.DataFrame(index=[3, 4], data={'a': [1.0, 2.0]}),
metadata=sentinel.meta)
assert not vs._delete_version.called
vs.write.assert_called_once_with(sentinel.symbol,
ANY,
prune_previous_version=True,
metadata=sentinel.meta)
vs.list_versions.assert_called_once_with(sentinel.symbol)
def test_ArcticTransaction_guards_against_inconsistent_ts():
vs = Mock(spec=VersionStore)
ts1 = pd.DataFrame(index=[1, 2], data={'a': [1.0, 2.0]})
vs.read.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=1,
metadata=None,
data=ts1,
host=sentinel.host)
vs.write.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=2,
metadata=None,
data=None,
host=sentinel.host)
vs.list_versions.side_effect = [{'version': 2}, {'version': 1}]
ts1 = pd.DataFrame(index=[1, 2], data={'a': [2.0, 3.0]})
with pytest.raises(ConcurrentModificationException):
with ArcticTransaction(vs,
sentinel.symbol,
sentinel.user,
sentinel.log,
modify_timeseries=ts1) as cwb:
pass
def test_ArcticTransaction_writes_no_data_found():
vs = Mock(spec=VersionStore)
ts1 = pd.DataFrame(index=[1, 2], data={'a': [1.0, 2.0]})
vs.read.side_effect = NoDataFoundException('no data')
vs.write.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=1,
metadata=None,
data=None)
vs.list_versions.side_effect = [
[],
[{
'version': 1
}],
]
with ArcticTransaction(vs, sentinel.symbol, sentinel.user,
sentinel.log) as cwb:
cwb.write(sentinel.symbol, ts1, metadata={1: 2})
assert vs.write.call_args_list == [
call(sentinel.symbol,
ANY,
prune_previous_version=True,
metadata={1: 2})
]
assert vs.list_versions.call_args_list == [
call(sentinel.symbol, latest_only=True),
call(sentinel.symbol)
]
def test_write_empty_metadata():
vs = _create_mock_versionstore()
expected_new_version = TPL_VERSION.copy()
expected_new_version.update({
'_id': MOCK_OBJID,
'version': TPL_VERSION['version'] + 1,
'metadata': None
})
expected_ret_val = VersionedItem(symbol=TEST_SYMBOL,
library=vs._arctic_lib.get_name(),
host=vs._arctic_lib.arctic.mongo_host,
version=TPL_VERSION['version'] + 1,
metadata=None,
data=None)
with patch('arctic.store.version_store.bson.ObjectId') as mock_objId, \
patch('arctic.store.version_store.mongo_retry') as mock_retry:
mock_objId.return_value = MOCK_OBJID
mock_retry.side_effect = lambda f: f
assert expected_ret_val == VersionStore.write_metadata(
vs, symbol=TEST_SYMBOL, metadata=None)
assert vs._versions.insert_one.call_args_list == [
call(expected_new_version)
]
assert vs._versions.delete_one.called is False
assert vs._publish_change.call_args_list == [
call(TEST_SYMBOL, expected_new_version)
]
assert vs.write.called is False
def _create_mock_versionstore():
vs = create_autospec(VersionStore,
_arctic_lib=Mock(),
_version_nums=Mock(),
_versions=Mock())
vs._insert_version = lambda version: VersionStore._insert_version(
vs, version)
vs._arctic_lib.get_name.return_value = TEST_LIB
vs._read_metadata.return_value = TPL_VERSION
vs._version_nums.find_one_and_update.return_value = {
'version': TPL_VERSION['version'] + 1
}
vs._version_nums.find_one.return_value = {
'version': TPL_VERSION['version'] + 1
}
vs._versions.find_one.return_value = TPL_VERSION
vs._add_new_version_using_reference.side_effect = lambda *args: VersionStore._add_new_version_using_reference(
vs, *args)
vs._last_version_seqnum = lambda version: VersionStore._last_version_seqnum(
vs, version)
vs.write.return_value = VersionedItem(
symbol=TEST_SYMBOL,
library=vs._arctic_lib.get_name(),
version=TPL_VERSION['version'] + 1,
metadata=META_TO_WRITE,
data=None,
host=vs._arctic_lib.arctic.mongo_host)
return vs
def test_restore_version():
vs = _create_mock_versionstore()
expected_new_version = TPL_VERSION.copy()
expected_new_version.update({
'_id': MOCK_OBJID,
'version': TPL_VERSION['version'] + 1,
'metadata': None
})
with patch('arctic.store.version_store.bson.ObjectId') as mock_objId, \
patch('arctic.store.version_store.mongo_retry') as mock_retry:
mock_objId.return_value = MOCK_OBJID
mock_retry.side_effect = lambda f: f
ret_val = VersionStore.restore_version(vs,
symbol=TEST_SYMBOL,
as_of=TPL_VERSION['version'],
prune_previous_version=True)
assert ret_val == VersionedItem(symbol=TEST_SYMBOL,
library=vs._arctic_lib.get_name(),
version=TPL_VERSION['version'] + 1,
metadata=None,
data=None)
assert vs._versions.insert_one.call_args_list == [
call(expected_new_version)
]
assert vs._publish_change.call_args_list == [
call(TEST_SYMBOL, expected_new_version)
]
def test_ArcticTransaction_writes_if_metadata_changed():
vs = Mock(spec=VersionStore)
ts1 = pd.DataFrame(index=[1, 2], data={'a': [1.0, 2.0]})
vs.read.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=1,
metadata=None,
data=ts1,
host=sentinel.host)
vs.write.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=2,
metadata=None,
data=None,
host=sentinel.host)
vs.list_versions.return_value = [{'version': 2}, {'version': 1}]
with ArcticTransaction(vs, sentinel.symbol, sentinel.user,
sentinel.log) as cwb:
assert cwb._do_write is False
cwb.write(sentinel.symbol, ts1, metadata={1: 2})
assert cwb._do_write is True
assert not vs._delete_version.called
vs.write.assert_called_once_with(sentinel.symbol,
ANY,
prune_previous_version=True,
metadata={1: 2})
vs.list_versions.assert_called_once_with(sentinel.symbol)
# Won't write on exit with same data and metadata
vs.read.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=2,
metadata={1: 2},
data=ts1,
host=sentinel.host)
with ArcticTransaction(vs, sentinel.symbol, sentinel.user,
sentinel.log) as cwb:
assert cwb._do_write is False
cwb.write(sentinel.symbol, ts1, metadata={1: 2})
assert cwb._do_write is False
def test_ArcticTransaction_does_nothing_when_data_is_None():
vs = Mock(spec=VersionStore)
ts1 = pd.DataFrame(index=[1, 2], data={'a': [1.0, 2.0]})
vs.read.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=1,
metadata=None,
data=ts1)
vs.write.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=2,
metadata=None,
data=None)
vs.list_versions.return_value = [{'version': 1}, {'version': 2}]
with ArcticTransaction(vs, sentinel.symbol, sentinel.user,
sentinel.log) as cwb:
pass
assert not vs._delete_version.called
assert not vs.write.called
def test_ArcticTransaction_does_nothing_when_data_not_modified():
vs = create_autospec(VersionStore, _collection=Mock())
ts1 = pd.DataFrame(index=[1, 2], data={'a': [1.0, 2.0]})
vs.read.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=1,
metadata=None,
data=ts1)
vs.write.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=2,
metadata=None,
data=None)
vs.list_versions.side_effect = [{'version': 2}, {'version': 1}]
with ArcticTransaction(vs, sentinel.symbol, sentinel.user,
sentinel.log) as cwb:
cwb.write(sentinel.symbol,
pd.DataFrame(index=[1, 2], data={'a': [1.0, 2.0]}))
assert not vs._delete_version.called
assert not vs.write.called
def test_ArcticTransaction_detects_concurrent_writes():
vs = Mock(spec=VersionStore)
ts1 = pd.DataFrame(index=[1, 2], data={'a': [1.0, 2.0]})
vs.read.return_value = VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=1,
metadata=None,
data=ts1)
vs.write.side_effect = [
VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=2,
metadata=None,
data=None),
VersionedItem(symbol=sentinel.symbol,
library=sentinel.library,
version=3,
metadata=None,
data=None)
]
#note that we return some extra version 5, it is possible that we have a write coming in after our own write that gets picked up
vs.list_versions.side_effect = [[
{
'version': 5
},
{
'version': 2
},
{
'version': 1
},
], [
{
'version': 5
},
{
'version': 3
},
{
'version': 2
},
{
'version': 1
},
]]
from threading import Event, Thread
e1 = Event()
e2 = Event()
def losing_writer():
#will attempt to write version 2, should find that version 2 is there and it ends up writing version 3
with pytest.raises(ArcticTransaction):
with ArcticTransaction(vs, sentinel.symbol, sentinel.user,
sentinel.log) as cwb:
cwb.write(sentinel.symbol, pd.DataFrame([1.0, 2.0], [3, 4]))
e1.wait()
def winning_writer():
#will attempt to write version 2 as well
with ArcticTransaction(vs, sentinel.symbol, sentinel.user,
sentinel.log) as cwb:
cwb.write(sentinel.symbol, pd.DataFrame([1.0, 2.0], [5, 6]))
e2.wait()
t1 = Thread(target=losing_writer)
t2 = Thread(target=winning_writer)
t1.start()
t2.start()
# both read the same timeseries and are locked doing some 'work'
e2.set()
# t2 should now be able to finish
t2.join()
e1.set()
t1.join()
# we're expecting the losing_writer to undo its write once it realises that it wrote v3 instead of v2
vs._delete_version.assert_called_once_with(sentinel.symbol, 3)
