#!/usr/bin/env python# -*- coding: utf-8 -*-## Copyright 2011-2013 Codernity (http://codernity.com)## Licensed under the Apache License, Version 2.0 (the "License");# you may not use this file except in compliance with the License.# You may obtain a copy of the License at## http://www.apache.org/licenses/LICENSE-2.0## Unless required by applicable law or agreed to in writing, software# distributed under the License is distributed on an "AS IS" BASIS,# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.# See the License for the specific language governing permissions and# limitations under the License.fromCodernityDB.databaseimportDatabase,RecordDeleted,RecordNotFoundfromCodernityDB.databaseimportDatabaseExceptionfromCodernityDB.hash_indeximportHashIndex,UniqueHashIndexfromCodernityDB.indeximportIndexExceptionfromCodernityDB.miscimportrandom_hex_32fromCodernityDBimportrr_cacheimportpytestimportosimportrandomfromhashlibimportmd5try:fromcollectionsimportCounterexceptImportError:classCounter(dict):'Mapping where default values are zero'def__missing__(self,key):return0classCustomHashIndex(HashIndex):def__init__(self,*args,**kwargs):kwargs['key_format']='I'kwargs['hash_lim']=1super(CustomHashIndex,self).__init__(*args,**kwargs)defmake_key_value(self,data):d=data.get('test')ifdisNone:returnNoneifd>5:k=1else:k=0returnk,dict(test=d)defmake_key(self,key):returnkeyclassMd5Index(HashIndex):def__init__(self,*args,**kwargs):kwargs['key_format']='16s'kwargs['hash_lim']=4*1024super(Md5Index,self).__init__(*args,**kwargs)defmake_key_value(self,data):returnmd5(data['name']).digest(),{}defmake_key(self,key):returnmd5(key).digest()classWithAIndex(HashIndex):def__init__(self,*args,**kwargs):kwargs['key_format']='16s'kwargs['hash_lim']=4*1024super(WithAIndex,self).__init__(*args,**kwargs)defmake_key_value(self,data):a_val=data.get("a")ifa_val:ifnotisinstance(a_val,basestring):a_val=str(a_val)returnmd5(a_val).digest(),{}returnNonedefmake_key(self,key):ifnotisinstance(key,basestring):key=str(key)returnmd5(key).digest()classHashIndexTests:defsetup_method(self,method):self.counter=Counter()deftest_simple(self,tmpdir):db=self._db(os.path.join(str(tmpdir),'db'))# db.set_indexes([UniqueHashIndex(db.path, 'id')])# db.initialize()db.add_index(UniqueHashIndex(db.path,'id'),False)db.create()withpytest.raises(RecordNotFound):db.get("id","not_existing")a1=dict(a=1)db.insert(a1)test=db.get('id',a1['_id'])asserttest['_id']==a1['_id']a_id=a1['_id']asserttest==a1a1['x']='x'db.update(a1)test2=db.get('id',a_id)asserttest2['x']=='x'asserttest2!=testdb.delete(a1)withpytest.raises(RecordDeleted):db.get('id',a_id)db.close()deftest_insert_with_id(self,tmpdir):db=self._db(os.path.join(str(tmpdir),'db'))# db.set_indexes([UniqueHashIndex(db.path, 'id')])# db.initialize()db.add_index(UniqueHashIndex(db.path,'id'),False)db.create()doc=dict(a=1,_id=random_hex_32())ins=db.insert(doc)assertins['_id']==doc["_id"]db.close()deftest_delete(self,tmpdir,inserts):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id')])db.create()ins=[]forxinxrange(inserts):doc=dict(x=x)db.insert(doc)ins.append(doc)self.counter['ins']+=1foriinxrange(inserts/10):curr=ins.pop(random.randint(0,len(ins)-1))db.delete(curr)assertlen(ins)==db.count(db.all,'id')forxinxrange(inserts):doc=dict(x=x)db.insert(doc)ins.append(doc)self.counter['ins']+=1foriinxrange(inserts/10):curr=ins.pop(random.randint(0,len(ins)-1))db.delete(curr)assertlen(ins)==db.count(db.all,'id')deftest_get_after_delete(self,tmpdir):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id'),CustomHashIndex(db.path,'custom')])db.create()forxinxrange(100):doc=dict(test=6)doc.update(db.insert(doc))ifdoc['test']>5:self.counter['r']+=1else:self.counter['l']+=1counted_bef=db.count(db.all,'custom')elem=db.all('custom',with_doc=True,limit=1)doc=next(elem)['doc']assertdb.delete(doc)==Truecounted_aft=db.count(db.all,'custom')assertcounted_bef-1==counted_aftfrom_ind=db.get('custom',1,with_doc=True)assertfrom_ind['doc']['_id']!=doc['_id']alls=db.all('custom',with_doc=True,limit=90)forcurrinalls:assertdb.delete(curr['doc'])==Truefrom_ind=db.get('custom',1,with_doc=True)assertfrom_ind['doc']!={}deftest_delete_with_id_and_rev_only(self,tmpdir,inserts):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id')])db.create()ins=[]forxinxrange(inserts):doc=dict(x=x)db.insert(doc)ins.append(doc)self.counter['ins']+=1foriinxrange(inserts/10):curr=ins.pop(random.randint(0,len(ins)-1))d={"_id":curr['_id'],'_rev':curr['_rev']}db.delete(d)assertlen(ins)==db.count(db.all,'id')forxinxrange(inserts):doc=dict(x=x)db.insert(doc)ins.append(doc)self.counter['ins']+=1foriinxrange(inserts/10):curr=ins.pop(random.randint(0,len(ins)-1))db.delete(curr)assertlen(ins)==db.count(db.all,'id')deftest_update_custom_unique(self,tmpdir,inserts):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id'),CustomHashIndex(db.path,'custom')])db.create()ins=[]forxinxrange(inserts):doc=dict(test=1)db.insert(doc)ins.append(doc)self.counter['ins']+=1assertlen(ins)==db.count(db.all,'id')assertlen(ins)==db.count(db.all,'custom')assertlen(ins)==db.count(db.get_many,'custom',key=0,limit=inserts+1)assert0==db.count(db.get_many,'custom',key=1,limit=inserts+1)sample=random.sample(ins,inserts/10)forcurrinsample:curr['test']=10db.update(curr)self.counter['upd']+=1assertself.counter['ins']==db.count(db.all,'id')assertself.counter['ins']==db.count(db.all,'custom')assertself.counter['upd']==db.count(db.get_many,'custom',key=1,limit=inserts+1)assertself.counter['ins']-self.counter['upd']==db.count(db.get_many,'custom',key=0,limit=inserts+1)deftest_update_custom_same_key_new_value(self,tmpdir,inserts):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id'),CustomHashIndex(db.path,'custom')])db.create()inserted=[]forxinxrange(inserts):inserted.append(db.insert(dict(test=x)))inserted[-1]['test']=xforelininserted[::20]:foriinxrange(4):curr=db.get('id',el['_id'],with_storage=True)assertel['test']==curr['test']el['test']+=random.randint(1,3)db.update(el)assertlen(inserted)==db.count(db.all,'custom')db.close()deftest_double_insert(self,tmpdir):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id')])db.create()a_id='54bee5c4628648b5a742379a1de89b2d'a1=dict(a=1,_id=a_id)db.insert(a1)a2=dict(a=2,_id=a_id)withpytest.raises(IndexException):db.insert(a2)deftest_adv1(self,tmpdir,inserts):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id')])db.create()l=[]foriinxrange(inserts):c=dict(i=i)db.insert(c)l.append(c)forjinrange(inserts):curr=l[j]c=db.get('id',curr['_id'])assertc['_id']==curr['_id']assertc['i']==jforiinxrange(inserts):curr=l[i]c=db.get("id",curr['_id'])c['update']=Truedb.update(c)forjinxrange(inserts):curr=l[i]c=db.get('id',curr['_id'])assertc['update']==Trueforjinxrange(inserts):curr=l[j]c=db.get('id',curr['_id'])assertdb.delete(c)==Trueforjinxrange(inserts):withpytest.raises(RecordDeleted):db.get('id',l[j]['_id'])db.close()deftest_all(self,tmpdir,inserts):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id')])db.create()l=[]foriinxrange(inserts):c=dict(i=i)db.insert(c)l.append(c)assertdb.count(db.all,'id')==insertsto_delete=random.randint(0,inserts-1)db.delete(l[to_delete])assertdb.count(db.all,'id')==inserts-1deftest_compact(self,tmpdir):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id')])db.create()l=[]foriinxrange(10):c=dict(i=i)db.insert(c)l.append(c)foriinxrange(10):curr=l[i]c=db.get("id",curr['_id'])c['update']=Truec.update(db.update(c))db.compact()forjinrange(10):curr=l[j]c=db.get('id',curr['_id'])assertc['_id']==curr['_id']assertc['i']==jdb.close()deftest_compact2(self,tmpdir):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id')])db.create()l=[]foriinxrange(10):ifi%2==0:c=dict(i=i,even=True)else:c=dict(i=i)c.update(db.insert(c))l.append(c)fori,currinenumerate(db.all('id')):ifi%2:db.delete(curr)db.compact()db.compact()db.compact()assertlen(list(db.all('id',with_doc=True)))==5forcurrindb.all('id'):db.delete(curr)assertlen(list(db.all('id',with_doc=True)))==0deftest_similar(self,tmpdir):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id'),Md5Index(db.path,'md5')])db.create()a=dict(name='pigmej')db.insert(a)db.get('md5','pigmej')withpytest.raises(RecordNotFound):db.get("md5",'pigme')deftest_custom_index(self,tmpdir):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id'),CustomHashIndex(db.path,'custom')])db.create()l_1=[]l_0=[]foriinxrange(10):c=dict(test=i)db.insert(c)ifi>5:l_1.append(c)else:l_0.append(c)i=0forcurrindb.all('custom',with_doc=False):i+=1asserti==len(l_1)+len(l_0)db.compact()gen=db.get_many('custom',key=1,limit=10,offset=0,with_doc=True)got=[]whileTrue:try:d=gen.next()exceptStopIteration:breakelse:got.append(d)assertlen(l_1)==len(got)fordocinl_1:doc['test']=0db.update(doc)gen=db.get_many('custom',key=0,limit=100,offset=0,with_doc=False)got=[]whileTrue:try:d=gen.next()exceptStopIteration:breakelse:got.append(d)assertlen(l_1)+len(l_0)==len(got)db.compact()gen=db.get_many('custom',key=0,limit=100,offset=0,with_doc=False)got=[]whileTrue:try:d=gen.next()exceptStopIteration:breakelse:got.append(d)assertlen(l_1)+len(l_0)==len(got)deftest_custom_index_2(self,tmpdir):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id'),WithAIndex(db.path,'with_a')])db.create()all_ins=[]foriinrange(10):curr=dict(something='other')ifi%2==0:curr['a']=str(i)all_ins.append(curr)db.insert(curr)l_0=len(list(db.all('id')))l_1=len(list(db.all('with_a')))assertl_1!=l_0all_a=list(db.all('with_a',with_doc=True))curr_a=all_a[0]['doc']delcurr_a['a']db.update(curr_a)l_2=len(list(db.all('with_a')))assertl_2+1==l_1curr_a=all_a[-1]['doc']db.delete(curr_a)l_3=len(list(db.all('with_a')))assertl_3+2==l_1deftest_insert_delete_compact_get_huge(self,tmpdir,inserts):inserts*=10db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id')])db.create()l=[]foriinxrange(inserts):c=dict(i=i)db.insert(c)l.append(c)assertdb.count(db.all,'id')==insertswill_delete=random.randint(0,inserts/15)forxinrange(will_delete):to_delete=random.randint(0,len(l)-1)db.delete(l.pop(to_delete))assertdb.count(db.all,'id')==inserts-will_deletedb.compact()assertdb.count(db.all,'id')==inserts-will_deletedeftest_all_same_keys(self,tmpdir,inserts):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id'),WithAIndex(db.path,'with_a')])db.create()l=0r=0z=0data=[]foriinxrange(inserts):a=random.randint(0,10)ifa>5:r+=1elifa==0:z+=1else:l+=1c=dict(a=a)db.insert(c)data.append(c)assertl+r+z==db.count(db.all,"id")assertl+r==db.count(db.all,"with_a")deftest_update_same(self,tmpdir,inserts):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id'),WithAIndex(db.path,'with_a'),CustomHashIndex(db.path,'custom')])db.create()inserted={}foriinxrange(inserts):a=random.randint(1,9)# a = 10ifa>5:self.counter['r']+=1else:self.counter['l']+=1c=dict(test=a)db.insert(c)inserted[c['_id']]=cdef_check():assertself.counter['l']==db.count(db.get_many,'custom',key=0,limit=inserts)assertself.counter['r']==db.count(db.get_many,'custom',key=1,limit=inserts)assertself.counter['r']+self.counter['l']==db.count(db.all,"custom")_check()d=inserted.values()[inserts/2]forcurr_uinrange(20):last_test=d['test']iflast_test>5:d['test']=1self.counter['r']-=1self.counter['l']+=1else:d['test']=6self.counter['l']-=1self.counter['r']+=1d['upd']=curr_udb.update(d)_check()num=inserts/10num=numifnum<20else20to_update=random.sample(inserted.values(),num)fordinto_update:forcurr_uinrange(10):last_test=d['test']iflast_test>5:d['test']=1self.counter['r']-=1self.counter['l']+=1else:d['test']=6self.counter['l']-=1self.counter['r']+=1d['upd']=curr_udb.update(d)_check()to_update=random.sample(inserted.values(),num)forcurr_uinrange(10):fordinto_update:last_test=d['test']iflast_test>5:d['test']=1self.counter['r']-=1self.counter['l']+=1else:d['test']=6self.counter['l']-=1self.counter['r']+=1d['upd']=curr_udb.update(d)_check()db.close()deftest_insert_on_deleted(self,tmpdir):db=self._db(os.path.join(str(tmpdir),'db'))db.create()a=db.insert(dict(_id='dd0604e2110d41e9a2ec630461ffd067'))db.insert(dict(_id='892508163deb4da0b44e8a00802dc75a'))db.delete(a)db.insert(dict(_id='d0a6e4e1aa74476a9012f9b8d7181a95'))db.get('id','892508163deb4da0b44e8a00802dc75a')db.close()deftest_offset_in_functions(self,tmpdir,inserts):db=self._db(os.path.join(str(tmpdir),'db'))db.set_indexes([UniqueHashIndex(db.path,'id'),CustomHashIndex(db.path,'custom')])db.create()offset=inserts//10or1real_inserts=insertsifinserts<1000else1000forxinxrange(real_inserts):db.insert(dict(test=x))assertreal_inserts-offset==db.count(db.all,'id',offset=offset)assertreal_inserts-offset==db.count(db.all,'custom',offset=offset)assert1==db.count(db.get_many,'custom',1,limit=1,offset=offset)db.close()