def _InitiateWebSocketConnection(self, local_conn, get_access_token_callback):
tunnel_target = self._GetTunnelTargetInfo()
new_websocket = iap_tunnel_websocket.IapTunnelWebSocket(
tunnel_target, get_access_token_callback,
functools.partial(_SendLocalDataCallback, local_conn),
functools.partial(_CloseLocalConnectionCallback, local_conn),
ignore_certs=self._ignore_certs)
new_websocket.InitiateConnection()
return new_websocket
def _InitiateWebSocketConnection(self, local_conn, cred):
tunnel_target = self._GetTunnelTargetInfo()
access_token = cred.access_token if cred else None
new_websocket = iap_tunnel_websocket.IapTunnelWebSocket(
tunnel_target, access_token,
functools.partial(self._SendLocalData, local_conn),
ignore_certs=self._ignore_certs)
new_websocket.InitiateConnection()
new_websocket.WaitForOpenOrRaiseError()
return new_websocket
def SetUp(self):
self._received_data = []
self._close_received = False
self.tunnel_target = utils.IapTunnelTargetInfo(
project='project-a', zone='zone-b', instance='instance-c',
interface='nic0', port=22, url_override=None, proxy_info=None)
self.iap_tunnel_websocket = iap_tunnel_websocket.IapTunnelWebSocket(
self.tunnel_target, GetAccessToken, self.DataCallback,
self.CloseCallback, ignore_certs=True)
self.iap_tunnel_websocket._websocket_helper = mock.MagicMock()
self.iap_tunnel_websocket._connect_msg_received = True
def testInit(self):
self.assertIs(self.iap_tunnel_websocket._tunnel_target, self.tunnel_target)
self.assertEqual(self.iap_tunnel_websocket._get_access_token_callback,
GetAccessToken)
self.iap_tunnel_websocket._data_handler_callback('test 1')
self.assertListEqual(self._received_data, ['test 1'])
self.assertTrue(self.iap_tunnel_websocket._ignore_certs)
second_socket = iap_tunnel_websocket.IapTunnelWebSocket(
self.tunnel_target, GetAccessToken, self.DataCallback,
self.CloseCallback)
self.assertFalse(second_socket._ignore_certs)
def testStartNewWebSocket(self, websocket_helper_cls_mock,
make_user_agent_mock):
new_socket = iap_tunnel_websocket.IapTunnelWebSocket(
self.tunnel_target, GetAccessToken, self.DataCallback,
self.CloseCallback, ignore_certs=True)
websocket_helper_mock = mock.MagicMock()
websocket_helper_cls_mock.return_value = websocket_helper_mock
make_user_agent_mock.return_value = 'my-user-agent/1234'
new_socket._StartNewWebSocket()
url = utils.CreateWebSocketConnectUrl(self.tunnel_target)
headers = ['User-Agent: my-user-agent/1234',
'Authorization: Bearer one-access-token']
websocket_helper_cls_mock.assert_called_once()
websocket_helper_cls_mock.assert_called_with(
url, headers, True, None, new_socket._OnData, new_socket._OnClose)
self.assertIs(new_socket._websocket_helper, websocket_helper_mock)
websocket_helper_mock.StartReceivingThread.assert_called_once()
