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()