mbligh | b62f724 | 2009-07-29 14:34:30 +0000 | [diff] [blame] | 1 | #!/usr/bin/python |
| 2 | # Copyright 2009 Google Inc. Released under the GPL v2 |
| 3 | |
| 4 | import unittest |
| 5 | |
| 6 | import common |
| 7 | from autotest_lib.mirror import database |
| 8 | from autotest_lib.client.common_lib.test_utils import mock |
| 9 | |
| 10 | class dict_database_unittest(unittest.TestCase): |
| 11 | _path = 'somepath.db' |
| 12 | |
| 13 | _db_contents = { |
| 14 | 'file1': database.item('file1', 10, 10000), |
| 15 | 'file2': database.item('file2', 20, 20000), |
| 16 | } |
| 17 | |
| 18 | class _file_instance: |
| 19 | pass |
| 20 | |
| 21 | def setUp(self): |
| 22 | self.god = mock.mock_god() |
| 23 | |
| 24 | self.god.stub_function(database.cPickle, 'load') |
| 25 | self.god.stub_function(database.cPickle, 'dump') |
| 26 | self.open_mock = self.god.create_mock_function('open') |
| 27 | |
| 28 | |
| 29 | def tearDown(self): |
| 30 | self.god.unstub_all() |
| 31 | |
| 32 | |
| 33 | def test_get_dictionary_no_file(self): |
| 34 | # record |
| 35 | (self.open_mock.expect_call(self._path, 'rb') |
| 36 | .and_raises(IOError('blah'))) |
| 37 | |
| 38 | # playback |
| 39 | db = database.dict_database(self._path) |
| 40 | self.assertEqual(db.get_dictionary(_open_func=self.open_mock), {}) |
| 41 | |
| 42 | self.god.check_playback() |
| 43 | |
| 44 | |
| 45 | def test_get_dictionary(self): |
| 46 | # record |
| 47 | (self.open_mock.expect_call(self._path, 'rb') |
| 48 | .and_return(self._file_instance)) |
| 49 | (database.cPickle.load.expect_call(self._file_instance) |
| 50 | .and_return(self._db_contents)) |
| 51 | |
| 52 | # playback |
| 53 | db = database.dict_database(self._path) |
| 54 | self.assertEqual(db.get_dictionary(_open_func=self.open_mock), |
| 55 | self._db_contents) |
| 56 | |
| 57 | self.god.check_playback() |
| 58 | |
| 59 | |
| 60 | def test_merge_dictionary(self): |
| 61 | # setup |
| 62 | db = database.dict_database(self._path) |
| 63 | self.god.stub_function(db, 'get_dictionary') |
| 64 | |
| 65 | new_files = { |
| 66 | 'file3': database.item('file3', 30, 30000), |
| 67 | 'file4': database.item('file4', 40, 40000), |
| 68 | } |
| 69 | all_files = dict(self._db_contents) |
| 70 | all_files.update(new_files) |
| 71 | |
| 72 | # record |
| 73 | db.get_dictionary.expect_call().and_return(self._db_contents) |
| 74 | |
| 75 | (self.open_mock.expect_call(self._path, 'wb') |
| 76 | .and_return(self._file_instance)) |
| 77 | database.cPickle.dump.expect_call(all_files, self._file_instance, |
| 78 | protocol=2) |
| 79 | |
| 80 | # playback |
| 81 | db.merge_dictionary(new_files, _open_func=self.open_mock) |
| 82 | self.god.check_playback() |
| 83 | |
| 84 | |
| 85 | if __name__ == '__main__': |
| 86 | unittest.main() |