Python version_store.VersionStore类代码示例

def test_write_metadata_insert_version_dupkeyerror():
    vs = _create_mock_versionstore()
    vs._versions.insert_one.__name__ = 'insert_one'
    vs._versions.insert_one.side_effect = [DuplicateKeyError('dup key error'), None]
    VersionStore.write_metadata(vs, symbol=TEST_SYMBOL, metadata=META_TO_WRITE)
    assert vs._version_nums.find_one_and_update.call_count == 2
    assert vs._versions.insert_one.call_count == 2

def test_write_metadata_insert_version_opfailure():
    vs = _create_mock_versionstore()
    vs._versions.insert_one.__name__ = 'insert_one'
    vs._versions.insert_one.side_effect = [OperationFailure('op failure'), None]
    VersionStore.write_metadata(vs, symbol=TEST_SYMBOL, metadata=META_TO_WRITE)
    assert vs._version_nums.find_one_and_update.call_count == 1
    assert vs._versions.insert_one.call_count == 2

def test_read_metadata_no_asof():
    # When we do a read, with naive as_of, that as_of is treated in London Time.
    vs = create_autospec(VersionStore, instance=True,
                     _versions=Mock(), _allow_secondary=False)
    VersionStore._read_metadata(vs, sentinel.symbol)
    versions = vs._versions.with_options.return_value
    assert versions.find_one.call_args_list == [call({'symbol': sentinel.symbol},
                                                         sort=[('version', pymongo.DESCENDING)])]

def test_initialize_library():
    arctic_lib = create_autospec(ArcticLibraryBinding)
    arctic_lib.arctic = create_autospec(Arctic, _allow_secondary=False)
    with patch('arctic.store.version_store.enable_sharding', autospec=True) as enable_sharding:
        arctic_lib.get_top_level_collection.return_value.database.create_collection.__name__ = 'some_name'
        arctic_lib.get_top_level_collection.return_value.database.collection_names.__name__ = 'some_name'
        VersionStore.initialize_library(arctic_lib, hashed=sentinel.hashed)
    assert enable_sharding.call_args_list == [call(arctic_lib.arctic, arctic_lib.get_name(), hashed=sentinel.hashed)]

def test_restore_version_insert_version_opfailure():
    vs = _create_mock_versionstore()
    vs._versions.insert_one.__name__ = 'insert_one'
    vs._versions.insert_one.side_effect = [OperationFailure('op failure'), None]
    VersionStore.restore_version(vs, symbol=TEST_SYMBOL,
                                 as_of=TPL_VERSION['version'], prune_previous_version=True)
    assert vs._version_nums.find_one_and_update.call_count == 1
    assert vs._versions.insert_one.call_count == 2
    assert vs._publish_change.call_count == 1

def test_read_as_of_NotNaive():
    # When we do a read, with naive as_of, that as_of is treated in London Time.
    vs = create_autospec(VersionStore, instance=True,
                     _versions=Mock(), _allow_secondary=False)
    VersionStore._read_metadata(vs, 'symbol', dt(2013, 4, 1, 9, 0, tzinfo=mktz('Europe/Paris')))
    versions = vs._versions.with_options.return_value
    versions.find_one.assert_called_once_with({'symbol':'symbol', '_id':
                                              {'$lt': bson.ObjectId.from_datetime(dt(2013, 4, 1, 9, 0, tzinfo=mktz('Europe/Paris')) + dtd(seconds=1))}},
                                             sort=[('_id', pymongo.DESCENDING)])

def test_restore_version_insert_version_dupkeyerror():
    vs = _create_mock_versionstore()
    vs._versions.insert_one.__name__ = 'insert_one'
    vs._versions.insert_one.side_effect = [DuplicateKeyError('dup key error'), None]
    VersionStore.restore_version(vs, symbol=TEST_SYMBOL,
                                 as_of=TPL_VERSION['version'], prune_previous_version=True)
    assert vs._version_nums.find_one_and_update.call_count == 2
    assert vs._versions.insert_one.call_count == 2
    assert vs._publish_change.call_count == 1

def test_restore_version_data_missing_symbol():
    vs = _create_mock_versionstore()
    vs._read_metadata.side_effect = NoDataFoundException("no data")
    with patch('arctic.store.version_store.mongo_retry') as mock_retry:
        mock_retry.side_effect = lambda f: f
        with pytest.raises(NoDataFoundException):
            VersionStore.restore_version(vs, symbol=TEST_SYMBOL,
                                         as_of=TPL_VERSION['version'], prune_previous_version=True)
    assert vs._read_metadata.call_args_list == [call(TEST_SYMBOL, as_of=TPL_VERSION['version'])]
    assert vs._versions.insert_one.called is False

def test_ensure_index():
    th = Mock()
    vs = create_autospec(VersionStore, _collection=Mock())
    with patch('arctic.store.version_store._TYPE_HANDLERS', [th]):
        VersionStore._ensure_index(vs)
    assert vs._collection.snapshots.create_index.call_args_list == [call([('name', 1)], unique=True, background=True)]
    assert vs._collection.versions.create_index.call_args_list == [call([('symbol', 1), ('_id', -1)], background=True),
                                                               call([('symbol', 1), ('version', -1)], unique=True, background=True)]
    assert vs._collection.version_nums.create_index.call_args_list == [call('symbol', unique=True, background=True)]
    th._ensure_index.assert_called_once_with(vs._collection)

def test_read_reports_random_errors():
    self = create_autospec(VersionStore, _versions=Mock(), _arctic_lib=Mock(),
                           _allow_secondary=True)
    self._collection = create_autospec(Collection)
    self._do_read.__name__ = 'name'  # feh: mongo_retry decorator cares about this
    self._do_read.side_effect = Exception('bad')
    with pytest.raises(Exception) as e:
        with patch('arctic.store.version_store.log_exception') as le:
            VersionStore.read(self, sentinel.symbol, sentinel.as_of, sentinel.from_version)
    assert 'bad' in str(e)
    assert le.call_count == 1

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._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.write.return_value = VersionedItem(symbol=TEST_SYMBOL, library=vs._arctic_lib.get_name(),
                                          version=TPL_VERSION['version'] + 1, metadata=META_TO_WRITE, data=None)
    return vs

def test_get_version_allow_secondary_True():
    vs = create_autospec(VersionStore, instance=True,
                         _versions=Mock())
    vs._allow_secondary = True
    vs._find_snapshots.return_value = 'snap'
    vs._versions.find.return_value = [{'_id': bson.ObjectId.from_datetime(dt(2013, 4, 1, 9, 0)),
                       'symbol': 's', 'version': 10}]

    VersionStore.read(vs, "symbol")
    assert vs._read_metadata.call_args_list == [call('symbol', as_of=None, read_preference=ReadPreference.NEAREST)]
    assert vs._do_read.call_args_list == [call('symbol', vs._read_metadata.return_value, None, read_preference=ReadPreference.NEAREST)]

def test_snapshot():
    vs = create_autospec(VersionStore, _snapshots=Mock(),
                                       _collection=Mock(),
                                       _versions=Mock())
    vs._snapshots.find_one.return_value = False
    vs._versions.update_one.__name__ = 'name'
    vs._snapshots.insert_one.__name__ = 'name'
    vs.list_symbols.return_value = ['foo', 'bar']
    VersionStore.snapshot(vs, "symbol")
    assert vs._read_metadata.call_args_list == [call('foo', as_of=None, read_preference=ReadPreference.PRIMARY),
                                                call('bar', as_of=None, read_preference=ReadPreference.PRIMARY)]

def test_write_ensure_index():
    write_handler = Mock(write=Mock(__name__=""))
    vs = create_autospec(VersionStore, instance=True,
                     _collection=Mock(),
                     _version_nums=Mock(find_one_and_update=Mock(return_value={'version':1})),
                     _versions=Mock(insert_one=lambda x:None),
                     _arctic_lib=Mock(),
                     _publish_changes=False)
    vs._collection.database.connection.nodes = []
    vs._write_handler.return_value = write_handler
    VersionStore.write(vs, 'sym', sentinel.data, prune_previous_version=False)
    vs._ensure_index.assert_called_once_with()

def test_get_version_allow_secondary_user_override_False():
    """Ensure user can override read preference when calling read"""
    vs = create_autospec(VersionStore, instance=True,
                         _versions=Mock())
    vs._allow_secondary = True
    vs._find_snapshots.return_value = 'snap'
    vs._versions.find.return_value = [{'_id': bson.ObjectId.from_datetime(dt(2013, 4, 1, 9, 0)),
                       'symbol': 's', 'version': 10}]

    VersionStore.read(vs, "symbol", allow_secondary=False)
    assert vs._read_metadata.call_args_list == [call('symbol', as_of=None, read_preference=ReadPreference.PRIMARY)]
    assert vs._do_read.call_args_list == [call('symbol', vs._read_metadata.return_value, None, read_preference=ReadPreference.PRIMARY)]