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 6 import mock 7 import unittest 8 9 import common 10 from autotest_lib.client.common_lib import error 11 from autotest_lib.server.hosts import jetstream_host 12 13 14 class 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 34 if __name__ == "__main__": 35 unittest.main() 36