| #!/usr/bin/python |
| # Copyright 2009 Google Inc. Released under the GPL v2 |
| |
| import unittest |
| |
| import common |
| from autotest_lib.mirror import database |
| from autotest_lib.client.common_lib.test_utils import mock |
| |
| class dict_database_unittest(unittest.TestCase): |
| _path = 'somepath.db' |
| |
| _db_contents = { |
| 'file1': database.item('file1', 10, 10000), |
| 'file2': database.item('file2', 20, 20000), |
| } |
| |
| class _file_instance: |
| pass |
| |
| def setUp(self): |
| self.god = mock.mock_god() |
| |
| self.god.stub_function(database.cPickle, 'load') |
| self.god.stub_function(database.cPickle, 'dump') |
| self.open_mock = self.god.create_mock_function('open') |
| |
| |
| def tearDown(self): |
| self.god.unstub_all() |
| |
| |
| def test_get_dictionary_no_file(self): |
| # record |
| (self.open_mock.expect_call(self._path, 'rb') |
| .and_raises(IOError('blah'))) |
| |
| # playback |
| db = database.dict_database(self._path) |
| self.assertEqual(db.get_dictionary(_open_func=self.open_mock), {}) |
| |
| self.god.check_playback() |
| |
| |
| def test_get_dictionary(self): |
| # record |
| (self.open_mock.expect_call(self._path, 'rb') |
| .and_return(self._file_instance)) |
| (database.cPickle.load.expect_call(self._file_instance) |
| .and_return(self._db_contents)) |
| |
| # playback |
| db = database.dict_database(self._path) |
| self.assertEqual(db.get_dictionary(_open_func=self.open_mock), |
| self._db_contents) |
| |
| self.god.check_playback() |
| |
| |
| def test_merge_dictionary(self): |
| # setup |
| db = database.dict_database(self._path) |
| self.god.stub_function(db, 'get_dictionary') |
| |
| new_files = { |
| 'file3': database.item('file3', 30, 30000), |
| 'file4': database.item('file4', 40, 40000), |
| } |
| all_files = dict(self._db_contents) |
| all_files.update(new_files) |
| |
| # record |
| db.get_dictionary.expect_call().and_return(self._db_contents) |
| |
| (self.open_mock.expect_call(self._path, 'wb') |
| .and_return(self._file_instance)) |
| database.cPickle.dump.expect_call(all_files, self._file_instance, |
| protocol=2) |
| |
| # playback |
| db.merge_dictionary(new_files, _open_func=self.open_mock) |
| self.god.check_playback() |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |