1#!/usr/bin/python 2# Copyright (c) 2017 The Chromium OS Authors. All rights reserved. 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5 6import mock 7import unittest 8 9import common 10from autotest_lib.client.common_lib import error 11from autotest_lib.server.hosts import jetstream_host 12 13 14class JetstreamHostTestCase(unittest.TestCase): 15 16 @mock.patch.object(jetstream_host.logging, 'exception') 17 @mock.patch.object(jetstream_host.JetstreamHost, 'cleanup_services') 18 def test_cleanup(self, mock_cleanup, mock_exception_logging): 19 host = jetstream_host.JetstreamHost('') 20 host.prepare_for_update() 21 mock_cleanup.assert_called_with() 22 mock_exception_logging.assert_not_called() 23 24 @mock.patch.object(jetstream_host.logging, 'exception') 25 @mock.patch.object(jetstream_host.JetstreamHost, 'cleanup_services') 26 def test_failed_cleanup(self, mock_cleanup, mock_exception_logging): 27 mock_cleanup.side_effect = error.AutoservRunError('failed', None) 28 host = jetstream_host.JetstreamHost('') 29 host.prepare_for_update() 30 mock_cleanup.assert_called_with() 31 mock_exception_logging.assert_called() 32 33 34if __name__ == "__main__": 35 unittest.main() 36