Python mock.create_autospec() Examples

def test_append_insert_version_duplicatekey():
    read_handler = Mock(append=Mock(__name__=""))
    previous_version = TPL_VERSION.copy()
    previous_version['version'] = 1
    vs = create_autospec(VersionStore, instance=True,
                         _collection=Mock(),
                         _version_nums=Mock(find_one_and_update=Mock(return_value={'version': previous_version['version']+1})),
                         _versions=Mock(insert_one=Mock(__name__="insert_one"), find_one=Mock(__name__="find_one", return_value=previous_version)),
                         _arctic_lib=create_autospec(ArcticLibraryBinding,
                                                     arctic=create_autospec(Arctic, mongo_host='some_host')))
    vs._insert_version = lambda version: VersionStore._insert_version(vs, version)
    vs._versions.insert_one.side_effect = [DuplicateKeyError("dup key error"), None]
    vs._collection.database.connection.nodes = []
    vs._read_handler.return_value = read_handler
    VersionStore.append(vs, 'sym', [1, 2, 3], prune_previous_version=False, upsert=False)
    assert vs._version_nums.find_one_and_update.call_count == 2
    assert vs._versions.find_one.call_count == 2
    assert read_handler.append.call_count == 2
    assert vs._versions.insert_one.call_count == 2 

def test_list_symbols_all_options():
    ms = create_autospec(MetadataStore)
    ms.aggregate.return_value = []

    expected_pipeline = [
        {'$sort': {'symbol': 1, 'start_time': -1}},
        {'$match': {'symbol': {'$regex': 'test.*'},
                    'start_time': {'$lte': dt.datetime(2018, 5, 11)}}},
        {'$group': {'_id': '$symbol', 'metadata': {'$first': '$metadata'}}},
        {'$match': {'metadata.foo': 'bar'}},
        {'$project': {'_id': 0, 'symbol':  '$_id'}}
    ]

    MetadataStore.list_symbols(ms,
                               regex='test.*',
                               as_of=dt.datetime(2018, 5, 11),
                               foo='bar')
    ms.aggregate.assert_called_once_with(expected_pipeline) 

def test_list_symbols_as_of_regex():
    ms = create_autospec(MetadataStore)
    ms.aggregate.return_value = []

    expected_pipeline = [
        {'$sort': {'symbol': 1, 'start_time': -1}},
        {'$match': {'symbol': {'$regex': 'test.*'},
                    'start_time': {'$lte': dt.datetime(2018, 5, 11)}}},
        {'$group': {'_id': '$symbol', 'metadata': {'$first': '$metadata'}}},
        {'$project': {'_id': 0, 'symbol':  '$_id'}}
    ]

    MetadataStore.list_symbols(ms,
                               regex='test.*',
                               as_of=dt.datetime(2018, 5, 11))
    ms.aggregate.assert_called_once_with(expected_pipeline) 

def test_list_versions_no_snapshot():
    vs = create_autospec(VersionStore, instance=True, _versions=Mock(), _snapshots=Mock())
    vs._snapshots.find.return_value = []
    vs._snapshots.find_one.return_value = []
    date = dt(2013, 4, 1, 9, 0)
    vs._versions.find.return_value = [{'_id': bson.ObjectId.from_datetime(date),
                                       'symbol': 's',
                                       'version': 10,
                                       'metadata': None,
                                       'parent': []}]
    version = list(VersionStore.list_versions(vs, "symbol"))[0]
    local_date = date.replace(tzinfo=mktz("UTC"))
    assert version == {'symbol': version['symbol'],
                       'version': version['version'],
                       'date': local_date,
                       'snapshots': [],
                       'deleted': False} 

def test_concat_and_rewrite_checks_written():
    self = create_autospec(NdarrayStore)
    collection = create_autospec(Collection)
    version = {'_id': sentinel.version_id,
               'segment_count': 1}
    previous_version = {'_id': sentinel.id,
                       'up_to': sentinel.up_to,
                        'base_version_id': sentinel.base_version_id,
                        'version': sentinel.version,
                        'segment_count' : 5,
                        'append_count' : 3}
    symbol = sentinel.symbol
    item = []

    collection.find.return_value = [{'_id': sentinel.id, 'segment': 47, 'compressed': True, 'sha': 'abc0'},
                                    {'_id': sentinel.id_2, 'segment': 48, 'compressed': True, 'sha': 'abc1'},
                                    # 3 appended items
                                    {'_id': sentinel.id_3, 'segment': 49, 'compressed': False, 'sha': 'abc2'},
                                    {'_id': sentinel.id_4, 'segment': 50, 'compressed': False, 'sha': 'abc3'},
                                    {'_id': sentinel.id_5, 'segment': 51, 'compressed': False, 'sha': 'abc4'}]
    collection.update_many.return_value = create_autospec(UpdateResult, matched_count=1)
    NdarrayStore._concat_and_rewrite(self, collection, version, symbol, item, previous_version)
    assert self.check_written.call_count == 1 

def test_check_quota_info():
    self = create_autospec(ArcticLibraryBinding, database_name='arctic_db',
                           library='lib')
    self.arctic = create_autospec(Arctic)
    self.get_library_metadata.return_value = 1024 * 1024 * 1024
    self.quota_countdown = 0
    self.arctic.__getitem__.return_value = Mock(stats=Mock(return_value={'totals':
                                                                             {'size': 1 * 1024 * 1024,
                                                                              'count': 100,
                                                                              }
                                                                             }))
    with patch('arctic.arctic.logger.info') as info:
        ArcticLibraryBinding.check_quota(self)
    self.arctic.__getitem__.assert_called_once_with(self.get_name.return_value)
    info.assert_called_once_with('Mongo Quota: arctic_db.lib 0.001 / 1 GB used')
    assert self.quota_countdown == 51153 

def test_concat_and_rewrite_checks_chunk_count():
    self = create_autospec(NdarrayStore)
    collection = create_autospec(Collection)
    version = {}
    previous_version = {'_id': sentinel.id,
                        'base_version_id': sentinel.base_version_id,
                        'version': sentinel.version,
                        'segment_count' : 3,
                        'append_count' : 1,
                        'up_to': sentinel.up_to}
    symbol = sentinel.symbol
    item = sentinel.item

    collection.find.return_value = [{'compressed': True, 'segment': 1},
                                    {'compressed': False, 'segment': 2}]
    with pytest.raises(DataIntegrityException) as e:
        NdarrayStore._concat_and_rewrite(self, collection, version, symbol, item, previous_version)
    assert str(e.value) == 'Symbol: sentinel.symbol:sentinel.version expected 1 segments but found 0' 

def test_ArcticLibraryBinding_db():
    arctic = create_autospec(Arctic)
    arctic._conn = create_autospec(MongoClient)
    alb = ArcticLibraryBinding(arctic, "sentinel.library")
    with patch.object(alb, '_auth') as _auth:
        # connection is cached during __init__
        alb._db
        assert _auth.call_count == 0

        # Change the arctic connection
        arctic._conn = create_autospec(MongoClient)
        alb._db
        assert _auth.call_count == 1

        # connection is still cached
        alb._db
        assert _auth.call_count == 1 

def test_append_insert_version_operror():
    read_handler = Mock(append=Mock(__name__=""))
    previous_version = TPL_VERSION.copy()
    previous_version['version'] = 1
    vs = create_autospec(VersionStore, instance=True,
                         _collection=Mock(),
                         _version_nums=Mock(find_one_and_update=Mock(return_value={'version': previous_version['version']+1})),
                         _versions=Mock(insert_one=Mock(__name__="insert_one"), find_one=Mock(__name__="find_one", return_value=previous_version)),
                         _arctic_lib=create_autospec(ArcticLibraryBinding,
                                                     arctic=create_autospec(Arctic, mongo_host='some_host')))
    vs._insert_version = lambda version: VersionStore._insert_version(vs, version)
    vs._versions.insert_one.side_effect = [OperationFailure("mongo op error"), None]
    vs._collection.database.connection.nodes = []
    vs._read_handler.return_value = read_handler
    VersionStore.append(vs, 'sym', [1, 2, 3], prune_previous_version=False, upsert=False)
    assert vs._version_nums.find_one_and_update.call_count == 1
    assert vs._versions.find_one.call_count == 1
    assert read_handler.append.call_count == 1
    assert vs._versions.insert_one.call_count == 2 

def test_write_insert_version_operror():
    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=Mock(__name__="insert_one"), find_one=Mock(__name__="find_one")),
                         _arctic_lib=create_autospec(ArcticLibraryBinding,
                                                     arctic=create_autospec(Arctic, mongo_host='some_host')))
    vs._insert_version = lambda version: VersionStore._insert_version(vs, version)
    vs._versions.insert_one.side_effect = [OperationFailure("mongo op error"), None]
    vs._collection.database.connection.nodes = []
    vs._write_handler.return_value = write_handler
    VersionStore.write(vs, 'sym', sentinel.data, prune_previous_version=False)
    assert vs._version_nums.find_one_and_update.call_count == 1
    assert vs._versions.find_one.call_count == 1
    assert write_handler.write.call_count == 1
    assert vs._versions.insert_one.call_count == 2 

def test_write_insert_version_duplicatekey():
    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=Mock(__name__="insert_one"), find_one=Mock(__name__="find_one")),
                         _arctic_lib=create_autospec(ArcticLibraryBinding,
                                                     arctic=create_autospec(Arctic, mongo_host='some_host')))
    vs._insert_version = lambda version: VersionStore._insert_version(vs, version)
    vs._versions.insert_one.side_effect = [DuplicateKeyError("dup key error"), None]
    vs._collection.database.connection.nodes = []
    vs._write_handler.return_value = write_handler
    VersionStore.write(vs, 'sym', sentinel.data, prune_previous_version=False)
    assert vs._version_nums.find_one_and_update.call_count == 2
    assert vs._versions.find_one.call_count == 2
    assert write_handler.write.call_count == 2
    assert vs._versions.insert_one.call_count == 2 

def test_write_error_clean_retry():
    write_handler = Mock(write=Mock(__name__=""))
    write_handler.write.side_effect = [OperationFailure("mongo failure"), None]
    vs = create_autospec(VersionStore, instance=True,
                         _collection=Mock(),
                         _version_nums=Mock(find_one_and_update=Mock(return_value={'version': 1})),
                         _versions=Mock(insert_one=Mock(__name__="insert_one"), find_one=Mock(__name__="find_one")),
                         _arctic_lib=create_autospec(ArcticLibraryBinding,
                                                     arctic=create_autospec(Arctic, mongo_host='some_host')))
    vs._insert_version = lambda version: VersionStore._insert_version(vs, version)
    vs._collection.database.connection.nodes = []
    vs._write_handler.return_value = write_handler
    VersionStore.write(vs, 'sym', sentinel.data, prune_previous_version=False)
    assert vs._version_nums.find_one_and_update.call_count == 2
    assert vs._versions.find_one.call_count == 2
    assert write_handler.write.call_count == 2
    assert vs._versions.insert_one.call_count == 1 

def test_prune_previous_versions_0_timeout():
    self = create_autospec(VersionStore, _versions=Mock())
    self.name = sentinel.name
    self._versions = create_autospec(Collection)
    self._versions.with_options.return_value.find.__name__ = 'find'
    self._versions.with_options.return_value.find.return_value = []
    with patch('arctic.store.version_store.dt') as dt:
        dt.utcnow.return_value = datetime.datetime(2013, 10, 1)
        VersionStore._find_prunable_version_ids(self, sentinel.symbol, keep_mins=0)
    assert self._versions.with_options.call_args_list == [call(read_preference=ReadPreference.PRIMARY)]
    assert self._versions.with_options.return_value.find.call_args_list == [
                                                  call({'$or': [{'parent': {'$exists': False}},
                                                                {'parent': []}],
                                                        'symbol': sentinel.symbol,
                                                        '_id': {'$lt': bson.ObjectId('524a10810000000000000000')}},
                                                       sort=[('version', -1)],
                                                       skip=1,
                                                       projection={'FW_POINTERS_CONFIG': 1, '_id': 1, 'SEGMENT_SHAS': 1}
                                                       )] 

def test_read_handles_operation_failure():
    self = Mock(spec=VersionStore)
    self._read_preference.return_value = sentinel.read_preference
    self._collection = create_autospec(Collection)
    self._read_metadata.side_effect = [sentinel.meta1, sentinel.meta2]
    self._read_metadata.__name__ = 'name'
    self._do_read.__name__ = 'name'  # feh: mongo_retry decorator cares about this
    self._do_read.side_effect = [OperationFailure('error'), sentinel.read]
    VersionStore.read(self, sentinel.symbol, sentinel.as_of,
                      from_version=sentinel.from_version,
                      date_range=sentinel.date_range,
                      other_kwarg=sentinel.other_kwarg)
    # Assert that, for the two read calls, the second uses the new metadata
    assert self._do_read.call_args_list == [call(sentinel.symbol, sentinel.meta1, 
                                                 sentinel.from_version,
                                                 date_range=sentinel.date_range,
                                                 other_kwarg=sentinel.other_kwarg,
                                                 read_preference=sentinel.read_preference)]
    assert self._do_read_retry.call_args_list == [call(sentinel.symbol, sentinel.meta2,
                                                       sentinel.from_version,
                                                       date_range=sentinel.date_range,
                                                       other_kwarg=sentinel.other_kwarg,
                                                       read_preference=ReadPreference.PRIMARY)] 

def test_read():
    self = create_autospec(TopLevelTickStore)
    tsl = TickStoreLibrary(create_autospec(TickStore), create_autospec(DateRange))
    self._get_libraries.return_value = [tsl, tsl]
    dr = create_autospec(DateRange)
    with patch('pandas.concat') as concat:
        res = TopLevelTickStore.read(self, sentinel.symbol, dr,
                                     columns=sentinel.include_columns,
                                     include_images=sentinel.include_images)
    assert concat.call_args_list == [call([tsl.library.read.return_value,
                                           tsl.library.read.return_value])]
    assert res == concat.return_value
    assert tsl.library.read.call_args_list == [call(sentinel.symbol, tsl.date_range.intersection.return_value,
                                                    sentinel.include_columns, include_images=sentinel.include_images),
                                               call(sentinel.symbol, tsl.date_range.intersection.return_value,
                                                    sentinel.include_columns, include_images=sentinel.include_images)] 

def test_write_pandas_data_to_right_libraries():
    self = create_autospec(TopLevelTickStore, _arctic_lib=MagicMock(), _collection=MagicMock())
    self._collection.find.return_value = [{'library_name': sentinel.libname1, 'start': sentinel.st1, 'end': sentinel.end1},
                                          {'library_name': sentinel.libname2, 'start': sentinel.st2, 'end': sentinel.end2}]
    slice1 = range(2)
    slice2 = range(4)
    when(self._slice).called_with(sentinel.data, sentinel.st1, sentinel.end1).then(slice1)
    when(self._slice).called_with(sentinel.data, sentinel.st2, sentinel.end2).then(slice2)
    mock_lib1 = Mock()
    mock_lib2 = Mock()
    when(self._arctic_lib.arctic.__getitem__).called_with(sentinel.libname1).then(mock_lib1)
    when(self._arctic_lib.arctic.__getitem__).called_with(sentinel.libname2).then(mock_lib2)
    with patch("arctic.tickstore.toplevel.to_dt") as patch_to_dt:
        patch_to_dt.side_effect = [sentinel.st1, sentinel.end1, sentinel.st2, sentinel.end2]
        TopLevelTickStore.write(self, 'blah', sentinel.data)
    mock_lib1.write.assert_called_once_with('blah', slice1)
    mock_lib2.write.assert_called_once_with('blah', slice2) 

def test_read_object_2():
    self = create_autospec(PickleStore)
    version = {'_id': sentinel._id,
               'blob': '__chunked__'}
    coll = Mock()
    arctic_lib = Mock()
    coll.find.return_value = [{'data': Binary(compressHC(cPickle.dumps(object))),
                               'symbol': 'sentinel.symbol',
                               'segment': 1}
                              ]
    arctic_lib.get_top_level_collection.return_value = coll

    assert PickleStore.read(self, arctic_lib, version, sentinel.symbol) == object
    assert coll.find.call_args_list == [call({'symbol': sentinel.symbol, 'parent': sentinel._id})] 

def test_delete_one():
    arctic_lib = create_autospec(ArcticLibraryBinding, instance=True)
    collection = create_autospec(Collection, instance=True)
    arctic_lib.get_top_level_collection.return_value = collection

    bsons = BSONStore(arctic_lib)
    bsons.delete_one(sentinel.filter)

    assert collection.delete_one.call_count == 1
    assert collection.delete_one.call_args_list == [call(sentinel.filter)] 

def test_find_one_and_delete():
    arctic_lib = create_autospec(ArcticLibraryBinding, instance=True)
    collection = create_autospec(Collection, instance=True)
    arctic_lib.get_top_level_collection.return_value = collection

    ms = BSONStore(arctic_lib)
    ms.find_one_and_delete(sentinel.filter)

    assert collection.find_one_and_delete.call_count == 1
    assert collection.find_one_and_delete.call_args_list == [call(sentinel.filter)] 

def test_bulk_write():
    arctic_lib = create_autospec(ArcticLibraryBinding, instance=True)
    collection = create_autospec(Collection, instance=True)
    arctic_lib.get_top_level_collection.return_value = collection

    bsons = BSONStore(arctic_lib)
    bsons.bulk_write(sentinel.requests)

    assert arctic_lib.check_quota.call_count == 1
    assert collection.bulk_write.call_count == 1
    assert collection.bulk_write.call_args_list == [call(sentinel.requests)] 

def test_distinct():
    arctic_lib = create_autospec(ArcticLibraryBinding, instance=True)
    collection = create_autospec(Collection, instance=True)
    arctic_lib.get_top_level_collection.return_value = collection

    bsons = BSONStore(arctic_lib)
    bsons.distinct(sentinel.key)

    assert collection.distinct.call_count == 1
    assert collection.distinct.call_args_list == [call(sentinel.key)] 

def test_ensure_index():
    ms = create_autospec(MetadataStore)
    MetadataStore._ensure_index(ms)
    assert ms.create_index.call_args_list == [call([('symbol', 1),
                                                    ('start_time', -1)],
                                                   unique=True,
                                                   background=True)] 

def test_add_observe_dt_index():
    self = create_autospec(BitemporalStore, observe_column='col_a')
    assert_frame_equal(BitemporalStore._add_observe_dt_index(self, ts1, as_of=dt(2001, 1, 1)),
                       read_str_as_pandas("""sample_dt |      col_a | near
                               2012-09-08 17:06:11.040 | 2001-01-01 |  1.0
                               2012-10-08 17:06:11.040 | 2001-01-01 |  2.0
                               2012-10-09 17:06:11.040 | 2001-01-01 |  2.5
                               2012-11-08 17:06:11.040 | 2001-01-01 |  3.0""", num_index=2)) 

def test_read_with_base_version_id():
    self = create_autospec(PickleStore)
    version = {'_id': sentinel._id,
               'base_version_id': sentinel.base_version_id,
               'blob': '__chunked__'}
    coll = Mock()
    arctic_lib = Mock()
    coll.find.return_value = [{'data': Binary(compressHC(cPickle.dumps(object))),
                               'symbol': 'sentinel.symbol',
                               'segment': 1}
                              ]
    arctic_lib.get_top_level_collection.return_value = coll

    assert PickleStore.read(self, arctic_lib, version, sentinel.symbol) == object
    assert coll.find.call_args_list == [call({'symbol': sentinel.symbol, 'parent': sentinel.base_version_id})] 

def test_snapshot_duplicate_raises_exception():
    vs = create_autospec(VersionStore, _snapshots=Mock())
    with pytest.raises(DuplicateSnapshotException) as e:
        vs._snapshots.find_one.return_value = True
        VersionStore.snapshot(vs, 'symbol')
        assert "Snapshot 'symbol' already exists" in str(e.value) 

def test_read():
    self = create_autospec(PickleStore)
    version = {'data': 'item'}
    assert PickleStore.read(self, sentinel.arctic_lib, version, sentinel.symbol) == 'item' 

def test_write_object():
    arctic_lib = Mock()
    self = create_autospec(PickleStore)
    version = {'_id': ObjectId()}
    PickleStore.write(self, arctic_lib, version, 'sentinel.symbol', sentinel.item, sentinel.previous_version)
    assert 'data' not in version

    assert version['blob'] == '__chunked__V2'
    coll = arctic_lib.get_top_level_collection.return_value
    assert coll.update_one.call_args_list == [call({'sha': checksum('sentinel.symbol', {'segment': 0, 'data': Binary(compress(cPickle.dumps(sentinel.item, cPickle.HIGHEST_PROTOCOL)))}),
                                                    'symbol': 'sentinel.symbol'},
                                                   {'$set': {'segment': 0, 'data': Binary(compress(cPickle.dumps(sentinel.item, cPickle.HIGHEST_PROTOCOL)), 0)},
                                                    '$addToSet': {'parent': version['_id']}}, upsert=True)] 

def test_write():
    self = create_autospec(PickleStore)
    version = {}
    PickleStore.write(self, sentinel.arctic_lib, version, sentinel.symbol, 'item', sentinel.previous_version)
    assert version['data'] == 'item' 

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_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)