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