1# Copyright (c) 2017 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Set of Mocks and Stubs for test SCPI & RF Switch.
6
7Implement a set of Mocks and Stubs which will be used in
8unittest for SCPI and RF Switch.
9"""
10
11import socket
12
13from autotest_lib.client.common_lib.test_utils import mock
14from autotest_lib.server.cros.network.rf_switch import rf_switch
15from autotest_lib.server.cros.network.rf_switch import scpi
16
17
18class SocketMock(mock.mock_class):
19    """Mock socket for scpi test."""
20
21    def __init__(self, cls, name, *args, **kwargs):
22        mock.mock_class.__init__(self, cls, name, *args, **kwargs)
23        self.host = None
24        self.port = None
25        self.recv_buffer = ''
26        self.send_data = ''
27
28    def close(self):
29        pass
30
31    def connect(self, host, port):
32        self.host = host
33        self.port = port
34
35    def send(self, data):
36        self.send_data = data
37
38    def recv(self, size):
39        if len(self.recv_val) > size:
40            return self.recv_val[:size]
41        return self.recv_val
42
43
44class ScpiMock(scpi.Scpi):
45    """Mock object for Scpi."""
46
47    def __init__(self, host='1.2.3.4', port=12345, test_interface=False):
48        if test_interface:
49            self.socket = SocketMock(socket, socket)
50            self.socket.connect(host, port)
51        else:
52            scpi.Scpi.__init__(self, host, port)
53
54
55class RfSwitchMock(rf_switch.RfSwitch):
56    """Mock object for RfSwitch."""
57
58    def __init__(self, host='1.2.3.4'):
59        self.socket = SocketMock(socket, socket)
60        self.socket.connect(host, scpi.Scpi.SCPI_PORT)
61