mirror/database_unittest.py

#!/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()