test_rocksdb.py (2626B)
1 # standard imports 2 import unittest 3 import os 4 import logging 5 import sys 6 import importlib 7 import tempfile 8 import shutil 9 10 # local imports 11 from shep.persist import PersistedState 12 from shep.error import ( 13 StateExists, 14 StateInvalid, 15 StateItemExists, 16 StateItemNotFound, 17 ) 18 19 logging.basicConfig(level=logging.DEBUG) 20 logg = logging.getLogger() 21 22 23 class TestRedisStore(unittest.TestCase): 24 25 def setUp(self): 26 from shep.store.rocksdb import RocksDbStoreFactory 27 self.d = tempfile.mkdtemp() 28 self.factory = RocksDbStoreFactory(self.d) 29 self.states = PersistedState(self.factory.add, 3) 30 self.states.add('foo') 31 self.states.add('bar') 32 self.states.add('baz') 33 34 35 def tearDown(self): 36 shutil.rmtree(self.d) 37 38 39 def test_add(self): 40 self.states.put('abcd', state=self.states.FOO, contents='baz') 41 v = self.states.get('abcd') 42 self.assertEqual(v, 'baz') 43 v = self.states.state('abcd') 44 self.assertEqual(v, self.states.FOO) 45 46 47 def test_next(self): 48 self.states.put('abcd') 49 50 self.states.next('abcd') 51 self.assertEqual(self.states.state('abcd'), self.states.FOO) 52 53 self.states.next('abcd') 54 self.assertEqual(self.states.state('abcd'), self.states.BAR) 55 56 self.states.next('abcd') 57 self.assertEqual(self.states.state('abcd'), self.states.BAZ) 58 59 with self.assertRaises(StateInvalid): 60 self.states.next('abcd') 61 62 v = self.states.state('abcd') 63 self.assertEqual(v, self.states.BAZ) 64 65 66 def test_replace(self): 67 with self.assertRaises(StateItemNotFound): 68 self.states.replace('abcd', contents='foo') 69 70 self.states.put('abcd', state=self.states.FOO, contents='baz') 71 self.states.replace('abcd', contents='bar') 72 v = self.states.get('abcd') 73 self.assertEqual(v, 'bar') 74 75 76 def test_factory_ls(self): 77 self.states.put('abcd') 78 self.states.put('xxxx', state=self.states.BAZ) 79 r = self.factory.ls() 80 self.assertEqual(len(r), 2) 81 82 self.states.put('yyyy', state=self.states.BAZ) 83 r = self.factory.ls() 84 self.assertEqual(len(r), 2) 85 86 self.states.put('zzzz', state=self.states.BAR) 87 r = self.factory.ls() 88 self.assertEqual(len(r), 3) 89 90 91 if __name__ == '__main__': 92 norocksdb = False 93 rocksdb = None 94 try: 95 importlib.import_module('rocksdb') 96 except ModuleNotFoundError: 97 logg.critical('rocksdb module not available, skipping tests.') 98 sys.exit(0) 99 100 unittest.main()