1"""
2TestCases for DB.associate.
3"""
4
5import sys, os, string
6import time
7from pprint import pprint
8
9import unittest
10from test_all import db, dbshelve, test_support, verbose, have_threads, \
11        get_new_environment_path
12
13
14#----------------------------------------------------------------------
15
16
17musicdata = {
181 : ("Bad English", "The Price Of Love", "Rock"),
192 : ("DNA featuring Suzanne Vega", "Tom's Diner", "Rock"),
203 : ("George Michael", "Praying For Time", "Rock"),
214 : ("Gloria Estefan", "Here We Are", "Rock"),
225 : ("Linda Ronstadt", "Don't Know Much", "Rock"),
236 : ("Michael Bolton", "How Am I Supposed To Live Without You", "Blues"),
247 : ("Paul Young", "Oh Girl", "Rock"),
258 : ("Paula Abdul", "Opposites Attract", "Rock"),
269 : ("Richard Marx", "Should've Known Better", "Rock"),
2710: ("Rod Stewart", "Forever Young", "Rock"),
2811: ("Roxette", "Dangerous", "Rock"),
2912: ("Sheena Easton", "The Lover In Me", "Rock"),
3013: ("Sinead O'Connor", "Nothing Compares 2 U", "Rock"),
3114: ("Stevie B.", "Because I Love You", "Rock"),
3215: ("Taylor Dayne", "Love Will Lead You Back", "Rock"),
3316: ("The Bangles", "Eternal Flame", "Rock"),
3417: ("Wilson Phillips", "Release Me", "Rock"),
3518: ("Billy Joel", "Blonde Over Blue", "Rock"),
3619: ("Billy Joel", "Famous Last Words", "Rock"),
3720: ("Billy Joel", "Lullabye (Goodnight, My Angel)", "Rock"),
3821: ("Billy Joel", "The River Of Dreams", "Rock"),
3922: ("Billy Joel", "Two Thousand Years", "Rock"),
4023: ("Janet Jackson", "Alright", "Rock"),
4124: ("Janet Jackson", "Black Cat", "Rock"),
4225: ("Janet Jackson", "Come Back To Me", "Rock"),
4326: ("Janet Jackson", "Escapade", "Rock"),
4427: ("Janet Jackson", "Love Will Never Do (Without You)", "Rock"),
4528: ("Janet Jackson", "Miss You Much", "Rock"),
4629: ("Janet Jackson", "Rhythm Nation", "Rock"),
4730: ("Janet Jackson", "State Of The World", "Rock"),
4831: ("Janet Jackson", "The Knowledge", "Rock"),
4932: ("Spyro Gyra", "End of Romanticism", "Jazz"),
5033: ("Spyro Gyra", "Heliopolis", "Jazz"),
5134: ("Spyro Gyra", "Jubilee", "Jazz"),
5235: ("Spyro Gyra", "Little Linda", "Jazz"),
5336: ("Spyro Gyra", "Morning Dance", "Jazz"),
5437: ("Spyro Gyra", "Song for Lorraine", "Jazz"),
5538: ("Yes", "Owner Of A Lonely Heart", "Rock"),
5639: ("Yes", "Rhythm Of Love", "Rock"),
5740: ("Cusco", "Dream Catcher", "New Age"),
5841: ("Cusco", "Geronimos Laughter", "New Age"),
5942: ("Cusco", "Ghost Dance", "New Age"),
6043: ("Blue Man Group", "Drumbone", "New Age"),
6144: ("Blue Man Group", "Endless Column", "New Age"),
6245: ("Blue Man Group", "Klein Mandelbrot", "New Age"),
6346: ("Kenny G", "Silhouette", "Jazz"),
6447: ("Sade", "Smooth Operator", "Jazz"),
6548: ("David Arkenstone", "Papillon (On The Wings Of The Butterfly)",
66     "New Age"),
6749: ("David Arkenstone", "Stepping Stars", "New Age"),
6850: ("David Arkenstone", "Carnation Lily Lily Rose", "New Age"),
6951: ("David Lanz", "Behind The Waterfall", "New Age"),
7052: ("David Lanz", "Cristofori's Dream", "New Age"),
7153: ("David Lanz", "Heartsounds", "New Age"),
7254: ("David Lanz", "Leaves on the Seine", "New Age"),
7399: ("unknown artist", "Unnamed song", "Unknown"),
74}
75
76#----------------------------------------------------------------------
77
78class AssociateErrorTestCase(unittest.TestCase):
79    def setUp(self):
80        self.filename = self.__class__.__name__ + '.db'
81        self.homeDir = get_new_environment_path()
82        self.env = db.DBEnv()
83        self.env.open(self.homeDir, db.DB_CREATE | db.DB_INIT_MPOOL)
84
85    def tearDown(self):
86        self.env.close()
87        self.env = None
88        test_support.rmtree(self.homeDir)
89
90    def test00_associateDBError(self):
91        if verbose:
92            print '\n', '-=' * 30
93            print "Running %s.test00_associateDBError..." % \
94                  self.__class__.__name__
95
96        dupDB = db.DB(self.env)
97        dupDB.set_flags(db.DB_DUP)
98        dupDB.open(self.filename, "primary", db.DB_BTREE, db.DB_CREATE)
99
100        secDB = db.DB(self.env)
101        secDB.open(self.filename, "secondary", db.DB_BTREE, db.DB_CREATE)
102
103        # dupDB has been configured to allow duplicates, it can't
104        # associate with a secondary.  Berkeley DB will return an error.
105        try:
106            def f(a,b): return a+b
107            dupDB.associate(secDB, f)
108        except db.DBError:
109            # good
110            secDB.close()
111            dupDB.close()
112        else:
113            secDB.close()
114            dupDB.close()
115            self.fail("DBError exception was expected")
116
117    @unittest.skipUnless(db.version() >= (4, 6), 'Needs 4.6+')
118    def test_associateListError(self):
119        db1 = db.DB(self.env)
120        db1.open('bad.db', "a.db", db.DB_BTREE, db.DB_CREATE)
121        db2 = db.DB(self.env)
122        db2.open('bad.db', "b.db", db.DB_BTREE, db.DB_CREATE)
123
124        db1.associate(db2, lambda a, b: [0])
125
126        msg = "TypeError: The list returned by DB->associate callback" \
127              " should be a list of strings."
128        with test_support.captured_output("stderr") as s:
129            db1.put("0", "1")
130        db1.close()
131        db2.close()
132        self.assertEquals(s.getvalue().strip(), msg)
133
134
135#----------------------------------------------------------------------
136
137
138class AssociateTestCase(unittest.TestCase):
139    keytype = ''
140    envFlags = 0
141    dbFlags = 0
142
143    def setUp(self):
144        self.filename = self.__class__.__name__ + '.db'
145        self.homeDir = get_new_environment_path()
146        self.env = db.DBEnv()
147        self.env.open(self.homeDir, db.DB_CREATE | db.DB_INIT_MPOOL |
148                               db.DB_INIT_LOCK | db.DB_THREAD | self.envFlags)
149
150    def tearDown(self):
151        self.closeDB()
152        self.env.close()
153        self.env = None
154        test_support.rmtree(self.homeDir)
155
156    def addDataToDB(self, d, txn=None):
157        for key, value in musicdata.items():
158            if type(self.keytype) == type(''):
159                key = "%02d" % key
160            d.put(key, '|'.join(value), txn=txn)
161
162    def createDB(self, txn=None):
163        self.cur = None
164        self.secDB = None
165        self.primary = db.DB(self.env)
166        self.primary.set_get_returns_none(2)
167        self.primary.open(self.filename, "primary", self.dbtype,
168                      db.DB_CREATE | db.DB_THREAD | self.dbFlags, txn=txn)
169
170    def closeDB(self):
171        if self.cur:
172            self.cur.close()
173            self.cur = None
174        if self.secDB:
175            self.secDB.close()
176            self.secDB = None
177        self.primary.close()
178        self.primary = None
179
180    def getDB(self):
181        return self.primary
182
183
184    def _associateWithDB(self, getGenre):
185        self.createDB()
186
187        self.secDB = db.DB(self.env)
188        self.secDB.set_flags(db.DB_DUP)
189        self.secDB.set_get_returns_none(2)
190        self.secDB.open(self.filename, "secondary", db.DB_BTREE,
191                   db.DB_CREATE | db.DB_THREAD | self.dbFlags)
192        self.getDB().associate(self.secDB, getGenre)
193
194        self.addDataToDB(self.getDB())
195
196        self.finish_test(self.secDB)
197
198    def test01_associateWithDB(self):
199        if verbose:
200            print '\n', '-=' * 30
201            print "Running %s.test01_associateWithDB..." % \
202                  self.__class__.__name__
203
204        return self._associateWithDB(self.getGenre)
205
206    def _associateAfterDB(self, getGenre) :
207        self.createDB()
208        self.addDataToDB(self.getDB())
209
210        self.secDB = db.DB(self.env)
211        self.secDB.set_flags(db.DB_DUP)
212        self.secDB.open(self.filename, "secondary", db.DB_BTREE,
213                   db.DB_CREATE | db.DB_THREAD | self.dbFlags)
214
215        # adding the DB_CREATE flag will cause it to index existing records
216        self.getDB().associate(self.secDB, getGenre, db.DB_CREATE)
217
218        self.finish_test(self.secDB)
219
220    def test02_associateAfterDB(self):
221        if verbose:
222            print '\n', '-=' * 30
223            print "Running %s.test02_associateAfterDB..." % \
224                  self.__class__.__name__
225
226        return self._associateAfterDB(self.getGenre)
227
228    if db.version() >= (4, 6):
229        def test03_associateWithDB(self):
230            if verbose:
231                print '\n', '-=' * 30
232                print "Running %s.test03_associateWithDB..." % \
233                      self.__class__.__name__
234
235            return self._associateWithDB(self.getGenreList)
236
237        def test04_associateAfterDB(self):
238            if verbose:
239                print '\n', '-=' * 30
240                print "Running %s.test04_associateAfterDB..." % \
241                      self.__class__.__name__
242
243            return self._associateAfterDB(self.getGenreList)
244
245
246    def finish_test(self, secDB, txn=None):
247        # 'Blues' should not be in the secondary database
248        vals = secDB.pget('Blues', txn=txn)
249        self.assertEqual(vals, None, vals)
250
251        vals = secDB.pget('Unknown', txn=txn)
252        self.assertIn(vals[0], (99, '99'), vals)
253        vals[1].index('Unknown')
254        vals[1].index('Unnamed')
255        vals[1].index('unknown')
256
257        if verbose:
258            print "Primary key traversal:"
259        self.cur = self.getDB().cursor(txn)
260        count = 0
261        rec = self.cur.first()
262        while rec is not None:
263            if type(self.keytype) == type(''):
264                self.assertTrue(int(rec[0]))  # for primary db, key is a number
265            else:
266                self.assertTrue(rec[0])
267                self.assertIs(type(rec[0]), int)
268            count = count + 1
269            if verbose:
270                print rec
271            rec = getattr(self.cur, "next")()
272        self.assertEqual(count, len(musicdata))  # all items accounted for
273
274
275        if verbose:
276            print "Secondary key traversal:"
277        self.cur = secDB.cursor(txn)
278        count = 0
279
280        # test cursor pget
281        vals = self.cur.pget('Unknown', flags=db.DB_LAST)
282        self.assertIn(vals[1], (99, '99'), vals)
283        self.assertEqual(vals[0], 'Unknown')
284        vals[2].index('Unknown')
285        vals[2].index('Unnamed')
286        vals[2].index('unknown')
287
288        vals = self.cur.pget('Unknown', data='wrong value', flags=db.DB_GET_BOTH)
289        self.assertEqual(vals, None, vals)
290
291        rec = self.cur.first()
292        self.assertEqual(rec[0], "Jazz")
293        while rec is not None:
294            count = count + 1
295            if verbose:
296                print rec
297            rec = getattr(self.cur, "next")()
298        # all items accounted for EXCEPT for 1 with "Blues" genre
299        self.assertEqual(count, len(musicdata)-1)
300
301        self.cur = None
302
303    def getGenre(self, priKey, priData):
304        self.assertEqual(type(priData), type(""))
305        genre = priData.split('|')[2]
306
307        if verbose:
308            print 'getGenre key: %r data: %r' % (priKey, priData)
309
310        if genre == 'Blues':
311            return db.DB_DONOTINDEX
312        else:
313            return genre
314
315    def getGenreList(self, priKey, PriData) :
316        v = self.getGenre(priKey, PriData)
317        if type(v) == type("") :
318            v = [v]
319        return v
320
321
322#----------------------------------------------------------------------
323
324
325class AssociateHashTestCase(AssociateTestCase):
326    dbtype = db.DB_HASH
327
328class AssociateBTreeTestCase(AssociateTestCase):
329    dbtype = db.DB_BTREE
330
331class AssociateRecnoTestCase(AssociateTestCase):
332    dbtype = db.DB_RECNO
333    keytype = 0
334
335#----------------------------------------------------------------------
336
337class AssociateBTreeTxnTestCase(AssociateBTreeTestCase):
338    envFlags = db.DB_INIT_TXN
339    dbFlags = 0
340
341    def txn_finish_test(self, sDB, txn):
342        try:
343            self.finish_test(sDB, txn=txn)
344        finally:
345            if self.cur:
346                self.cur.close()
347                self.cur = None
348            if txn:
349                txn.commit()
350
351    def test13_associate_in_transaction(self):
352        if verbose:
353            print '\n', '-=' * 30
354            print "Running %s.test13_associateAutoCommit..." % \
355                  self.__class__.__name__
356
357        txn = self.env.txn_begin()
358        try:
359            self.createDB(txn=txn)
360
361            self.secDB = db.DB(self.env)
362            self.secDB.set_flags(db.DB_DUP)
363            self.secDB.set_get_returns_none(2)
364            self.secDB.open(self.filename, "secondary", db.DB_BTREE,
365                       db.DB_CREATE | db.DB_THREAD, txn=txn)
366            self.getDB().associate(self.secDB, self.getGenre, txn=txn)
367
368            self.addDataToDB(self.getDB(), txn=txn)
369        except:
370            txn.abort()
371            raise
372
373        self.txn_finish_test(self.secDB, txn=txn)
374
375
376#----------------------------------------------------------------------
377
378class ShelveAssociateTestCase(AssociateTestCase):
379
380    def createDB(self):
381        self.primary = dbshelve.open(self.filename,
382                                     dbname="primary",
383                                     dbenv=self.env,
384                                     filetype=self.dbtype)
385
386    def addDataToDB(self, d):
387        for key, value in musicdata.items():
388            if type(self.keytype) == type(''):
389                key = "%02d" % key
390            d.put(key, value)    # save the value as is this time
391
392
393    def getGenre(self, priKey, priData):
394        self.assertEqual(type(priData), type(()))
395        if verbose:
396            print 'getGenre key: %r data: %r' % (priKey, priData)
397        genre = priData[2]
398        if genre == 'Blues':
399            return db.DB_DONOTINDEX
400        else:
401            return genre
402
403
404class ShelveAssociateHashTestCase(ShelveAssociateTestCase):
405    dbtype = db.DB_HASH
406
407class ShelveAssociateBTreeTestCase(ShelveAssociateTestCase):
408    dbtype = db.DB_BTREE
409
410class ShelveAssociateRecnoTestCase(ShelveAssociateTestCase):
411    dbtype = db.DB_RECNO
412    keytype = 0
413
414
415#----------------------------------------------------------------------
416
417class ThreadedAssociateTestCase(AssociateTestCase):
418
419    def addDataToDB(self, d):
420        t1 = Thread(target = self.writer1,
421                    args = (d, ))
422        t2 = Thread(target = self.writer2,
423                    args = (d, ))
424
425        t1.setDaemon(True)
426        t2.setDaemon(True)
427        t1.start()
428        t2.start()
429        t1.join()
430        t2.join()
431
432    def writer1(self, d):
433        for key, value in musicdata.items():
434            if type(self.keytype) == type(''):
435                key = "%02d" % key
436            d.put(key, '|'.join(value))
437
438    def writer2(self, d):
439        for x in range(100, 600):
440            key = 'z%2d' % x
441            value = [key] * 4
442            d.put(key, '|'.join(value))
443
444
445class ThreadedAssociateHashTestCase(ShelveAssociateTestCase):
446    dbtype = db.DB_HASH
447
448class ThreadedAssociateBTreeTestCase(ShelveAssociateTestCase):
449    dbtype = db.DB_BTREE
450
451class ThreadedAssociateRecnoTestCase(ShelveAssociateTestCase):
452    dbtype = db.DB_RECNO
453    keytype = 0
454
455
456#----------------------------------------------------------------------
457
458def test_suite():
459    suite = unittest.TestSuite()
460
461    suite.addTest(unittest.makeSuite(AssociateErrorTestCase))
462
463    suite.addTest(unittest.makeSuite(AssociateHashTestCase))
464    suite.addTest(unittest.makeSuite(AssociateBTreeTestCase))
465    suite.addTest(unittest.makeSuite(AssociateRecnoTestCase))
466
467    suite.addTest(unittest.makeSuite(AssociateBTreeTxnTestCase))
468
469    suite.addTest(unittest.makeSuite(ShelveAssociateHashTestCase))
470    suite.addTest(unittest.makeSuite(ShelveAssociateBTreeTestCase))
471    suite.addTest(unittest.makeSuite(ShelveAssociateRecnoTestCase))
472
473    if have_threads:
474        suite.addTest(unittest.makeSuite(ThreadedAssociateHashTestCase))
475        suite.addTest(unittest.makeSuite(ThreadedAssociateBTreeTestCase))
476        suite.addTest(unittest.makeSuite(ThreadedAssociateRecnoTestCase))
477
478    return suite
479
480
481if __name__ == '__main__':
482    unittest.main(defaultTest='test_suite')
483