1#!/usr/bin/env python
2#
3# Copyright (C) 2017 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18import unittest
19
20try:
21    from unittest import mock
22except ImportError:
23    import mock
24
25from host_controller import invocation_thread
26from host_controller.tfc import command_attempt
27from host_controller.tradefed import remote_operation
28
29
30class InvocationThreadTest(unittest.TestCase):
31    """A test for invocation_thread.InvocationThread.
32
33    Attributes:
34        _remote_client: A mock remote_client.RemoteClient.
35        _tfc_client: A mock tfc_client.TfcClient.
36        _inv_thread: The InvocationThread being tested.
37    """
38
39    def setUp(self):
40        """Creates the InvocationThread."""
41        self._remote_client = mock.Mock()
42        self._tfc_client = mock.Mock()
43        attempt = command_attempt.CommandAttempt(
44                task_id="321-0",
45                attempt_id="abcd-1234",
46                hostname="host0",
47                device_serial="ABCDEF")
48        command = ["vts", "-m", "SampleShellTest"]
49        serials = ["serial123", "serial456"]
50        self._inv_thread = invocation_thread.InvocationThread(
51                self._remote_client, self._tfc_client,
52                attempt, command, serials)
53
54    def _GetSubmittedEventTypes(self):
55        """Gets the types of the events submitted by the mock TfcClient.
56
57        Returns:
58            A list of strings, the event types.
59        """
60        event_types = []
61        for args, kwargs in self._tfc_client.SubmitCommandEvents.call_args_list:
62            event_types.extend(event["type"] for event in args[0])
63        return event_types
64
65    def _GetSentOperationTypes(self):
66        """Gets the types of the operations sent by the mock RemoteClient.
67
68        Returns:
69            A list of strings, the operation types.
70        """
71        operation_types = [args[0].type for args, kwargs in
72                           self._remote_client.SendOperation.call_args_list]
73        return operation_types
74
75    def testAllocationFailed(self):
76        """Tests AllocationFailed event."""
77        self._remote_client.SendOperation.side_effect = (
78                lambda op: _RaiseExceptionForOperation(op, "ALLOCATE_DEVICE"))
79        self._inv_thread.run()
80        self.assertEqual([command_attempt.EventType.ALLOCATION_FAILED],
81                         self._GetSubmittedEventTypes())
82        self.assertEqual(["ALLOCATE_DEVICE"],
83                         self._GetSentOperationTypes())
84
85    def testExecuteFailed(self):
86        """Tests ExecuteFailed event."""
87        self._remote_client.SendOperation.side_effect = (
88                lambda op: _RaiseExceptionForOperation(op, "EXEC_COMMAND"))
89        self._inv_thread.run()
90        self.assertEqual([command_attempt.EventType.EXECUTE_FAILED],
91                         self._GetSubmittedEventTypes())
92        self.assertEqual(["ALLOCATE_DEVICE",
93                          "ALLOCATE_DEVICE",
94                          "EXEC_COMMAND",
95                          "FREE_DEVICE",
96                          "FREE_DEVICE"],
97                         self._GetSentOperationTypes())
98
99    def testConfigurationError(self):
100        """Tests ConfigurationError event."""
101        self._remote_client.SendOperation.side_effect = (
102                lambda op: _RaiseExceptionForOperation(op, "EXEC_COMMAND",
103                                                       "Config error: test"))
104        self._inv_thread.run()
105        self.assertEqual([command_attempt.EventType.CONFIGURATION_ERROR],
106                         self._GetSubmittedEventTypes())
107        self.assertEqual(["ALLOCATE_DEVICE",
108                          "ALLOCATE_DEVICE",
109                          "EXEC_COMMAND",
110                          "FREE_DEVICE",
111                          "FREE_DEVICE"],
112                         self._GetSentOperationTypes())
113
114    def testInvocationCompleted(self):
115        """Tests InvocationCompleted event."""
116        self._remote_client.WaitForCommandResult.side_effect = (
117                None, {"status": "INVOCATION_SUCCESS"})
118        self._inv_thread.run()
119        self.assertEqual([command_attempt.EventType.INVOCATION_STARTED,
120                          command_attempt.EventType.TEST_RUN_IN_PROGRESS,
121                          command_attempt.EventType.INVOCATION_COMPLETED],
122                         self._GetSubmittedEventTypes())
123        # GET_LAST_COMMAND_RESULT isn't called in mock WaitForCommandResult.
124        self.assertEqual(["ALLOCATE_DEVICE",
125                          "ALLOCATE_DEVICE",
126                          "EXEC_COMMAND",
127                          "FREE_DEVICE",
128                          "FREE_DEVICE"],
129                         self._GetSentOperationTypes())
130
131
132def _RaiseExceptionForOperation(operation, op_type, error_msg="unit test"):
133    """Raises exception for specific operation type.
134
135    Args:
136        operation: A remote_operation.RemoteOperation object.
137        op_type: A string, the expected type.
138        error_msg: The message in the exception.
139
140    Raises:
141        RemoteOperationException if the operation's type matches op_type.
142    """
143    if operation.type == op_type:
144        raise remote_operation.RemoteOperationException(error_msg)
145
146
147if __name__ == "__main__":
148    unittest.main()
149