Python tornado.websocket.websocket_connect() Examples

def test_auth_url_provided_forwards_cookies(self):
    auth_url = self.get_server_base_url() + '/fake-auth'
    encoded_query_args = urlparse.urlencode(
        {'jupyter_http_over_ws_auth_url': auth_url})

    request = self.get_ws_connection_request('http_over_websocket?' +
                                             encoded_query_args)

    client = yield websocket.websocket_connect(request)
    client.write_message(
        self.get_request_json('/requires-cookies-from-fake-auth', '1234'))

    response_body = yield client.read_message()
    response = json.loads(response_body)
    self.assertEqual(200, response['status'])
    # RequiresCookieFromFakeAuthHandler writes cookies that should have
    # been received by performing a request to the auth url.
    self.assertEqual('X-Test-Cookie: "1" X-Other-Test-Cookie: "2"',
                     base64.b64decode(response['data']).decode('utf-8'))
    self.assertEqual('1234', response['message_id'])
    self.assertTrue(response['done']) 

def _assertCrossDomainRequestFails(self, auth_url):
    encoded_query_args = urlparse.urlencode(
        {'jupyter_http_over_ws_auth_url': auth_url})

    request = self.get_ws_connection_request('http_over_websocket?' +
                                             encoded_query_args)

    with testing.ExpectLog(
        'tornado.application',
        'Uncaught error when proxying request',
        required=True) as expect_log:
      client = yield websocket.websocket_connect(request)
      client.write_message(
          self.get_request_json('/requires-cookies-from-fake-auth-ws', '1234'))

      response_body = yield client.read_message()
      response = json.loads(response_body)
      self.assertEqual('1234', response['message_id'])
      self.assertEqual(500, response['status'])
      self.assertEqual(
          'Uncaught server-side exception. Check logs for additional details.',
          response['statusText'])
      self.assertTrue(response['done'])
      self.assertTrue(expect_log.logged_stack) 

def test_diagnostic_handler_no_problems_request(self):
    request = self.get_ws_connection_request('http_over_websocket/diagnose')
    request.url += '?min_version=0.0.7'
    request.headers.add('Origin', ALLOWED_ORIGIN)
    request.headers.add('Cookie', '_xsrf=' + FAKE_XSRF_VALUE)

    client = yield websocket.websocket_connect(request)
    self.assertIsNotNone(client)
    client.write_message('1')

    response_body = yield client.read_message()
    response = json.loads(response_body)

    self.assertEqual(
        {
            'message_id': '1',
            'extension_version': '0.0.7',
            'has_authentication_cookie': True,
            'is_outdated_extension': False
        }, response) 

def test_diagnostic_handler_missing_xsrf_cookie(self):
    request = self.get_ws_connection_request('http_over_websocket/diagnose')
    request.headers.add('Origin', ALLOWED_ORIGIN)

    client = yield websocket.websocket_connect(request)
    client.write_message('1')

    response_body = yield client.read_message()
    response = json.loads(response_body)

    self.assertEqual(
        {
            'message_id': '1',
            'extension_version': '0.0.7',
            'has_authentication_cookie': False,
            'is_outdated_extension': False
        }, response) 

def test_diagnostic_handler_newer_protocol_version_requested(self):
    request = self.get_ws_connection_request('http_over_websocket/diagnose')
    request.url += '?min_version=0.0.8'
    request.headers.add('Origin', ALLOWED_ORIGIN)
    request.headers.add('Cookie', '_xsrf=' + FAKE_XSRF_VALUE)

    client = yield websocket.websocket_connect(request)
    client.write_message('1')

    response_body = yield client.read_message()
    response = json.loads(response_body)

    self.assertEqual(
        {
            'message_id': '1',
            'extension_version': '0.0.7',
            'has_authentication_cookie': True,
            'is_outdated_extension': True
        }, response) 

def _connect(self):
        ws = await websocket.websocket_connect(self._server_ws_url)
        ws.__class__ = SafeWebSocket

        await ws.write_message({
            "command": "handshake",
            "name": self._name,
            "owner": self._owner,
            "secret": self._secret,
            "url": self._provider_url,
            "priority": self._priority,  # the large the importanter
        })

        msg = await ws.read_message()
        logger.info("WS receive: %s", msg)
        return ws 

def test_auth_url_provided_fails(self):
    auth_url = self.get_server_base_url() + '/fake-auth?always_fail=1'
    encoded_query_args = urlparse.urlencode(
        {'jupyter_http_over_ws_auth_url': auth_url})

    request = self.get_ws_connection_request(
        'http_over_websocket/proxied_ws/requires-cookies-from-fake-auth-ws?' +
        encoded_query_args)
    request.headers.add('Origin', ALLOWED_ORIGIN)

    client = yield websocket.websocket_connect(request)
    self.assertIsNotNone(client)

    msg = yield client.read_message()
    self.assertIsNone(msg)
    self.assertEqual(500, client.close_code) 

def _assertCrossDomainRequestFails(self, auth_url):
    encoded_query_args = urlparse.urlencode(
        {'jupyter_http_over_ws_auth_url': auth_url})

    request = self.get_ws_connection_request(
        'http_over_websocket/proxied_ws/requires-cookies-from-fake-auth-ws?' +
        encoded_query_args)
    request.headers.add('Origin', ALLOWED_ORIGIN)

    with testing.ExpectLog(
        'tornado.application',
        'Uncaught error when proxying request',
        required=True) as expect_log:
      client = yield websocket.websocket_connect(request)
      self.assertIsNotNone(client)

      msg = yield client.read_message()
      # Message of None indicates that the connection has been closed.
      self.assertIsNone(msg)
      self.assertEqual(500, client.close_code)
      self.assertTrue(expect_log.logged_stack) 

def test_proxied_connection_io(self):
    request = self.get_ws_connection_request(
        'http_over_websocket/proxied_ws/echoing-ws?someparam=1')
    request.headers.add('Origin', ALLOWED_ORIGIN)
    request.headers.add('X-Echo-Header', 'example')

    client = yield websocket.websocket_connect(request)
    self.assertIsNotNone(client)

    # EchoingWebSocketHandler writes a message on connection with request
    # details.
    initial_message = yield client.read_message()
    self.assertEqual(
        '<ECHO INIT> GET /echoing-ws?someparam=1 X-Echo-Header: example',
        initial_message)

    client.write_message('1')
    echo_message = yield client.read_message()
    self.assertEqual('<ECHO RESPONSE> 1', echo_message)

    client.write_message('something else!')
    echo_message = yield client.read_message()
    self.assertEqual('<ECHO RESPONSE> something else!', echo_message) 

def test_auth_url_provided_fails(self):
    auth_url = self.get_server_base_url() + '/fake-auth?always_fail=1'
    encoded_query_args = urlparse.urlencode(
        {'jupyter_http_over_ws_auth_url': auth_url})

    request = self.get_ws_connection_request('http_over_websocket?' +
                                             encoded_query_args)

    client = yield websocket.websocket_connect(request)
    client.write_message(
        self.get_request_json('/requires-cookies-from-fake-auth', '1234'))

    response_body = yield client.read_message()
    response = json.loads(response_body)
    self.assertEqual('1234', response['message_id'])
    self.assertEqual(500, response['status'])
    self.assertEqual(
        'Uncaught server-side exception. Check logs for additional details.',
        response['statusText'])
    self.assertTrue(response['done']) 

def test_auth_url_provided_forwards_cookies_and_propagates_xsrf(self):
    auth_url = self.get_server_base_url() + '/fake-auth'
    encoded_query_args = urlparse.urlencode(
        {'jupyter_http_over_ws_auth_url': auth_url})

    request = self.get_ws_connection_request('http_over_websocket?' +
                                             encoded_query_args)
    client = yield websocket.websocket_connect(request)
    client.write_message(
        self.get_request_json(
            '/api/sessions', '1234', method='POST', body='somedata'))

    response_body = yield client.read_message()
    response = json.loads(response_body)
    self.assertEqual(200, response['status'])
    self.assertEqual('somedata',
                     base64.b64decode(response['data']).decode('utf-8'))
    self.assertEqual('1234', response['message_id'])
    self.assertIn(TEST_XSRF_HEADER, response['headers'])
    self.assertEqual(FAKE_XSRF_VALUE, response['headers'][TEST_XSRF_HEADER])
    self.assertTrue(response['done']) 

def run_tests():
    url = options.url + '/getCaseCount'
    control_ws = yield websocket_connect(url, None)
    num_tests = int((yield control_ws.read_message()))
    logging.info('running %d cases', num_tests)
    msg = yield control_ws.read_message()
    assert msg is None

    for i in range(1, num_tests + 1):
        logging.info('running test case %d', i)
        url = options.url + '/runCase?case=%d&agent=%s' % (i, options.name)
        test_ws = yield websocket_connect(url, None, compression_options={})
        while True:
            message = yield test_ws.read_message()
            if message is None:
                break
            test_ws.write_message(message, binary=isinstance(message, bytes))

    url = options.url + '/updateReports?agent=%s' % options.name
    update_ws = yield websocket_connect(url, None)
    msg = yield update_ws.read_message()
    assert msg is None
    IOLoop.instance().stop() 

def _connect(self):
        ws = await websocket.websocket_connect(self._ws_url, ping_interval=3)
        ws.__class__ = SafeWebSocket

        await ws.write_message({
            "command": "handshake",
            "name": self._name,
            "owner": self._owner,
            "secret": self._secret,
            "url": self._provider_url,
            "priority": self._priority,  # the large the importanter
        })

        msg = await ws.read_message()
        logger.info("WS receive: %s", msg)
        return ws 

def wsquery(wsurl,msg): # Synchrone call, with tornado
    """ In a simple world, could be (with websocket_client pypi):
            ws = websocket.create_connection(wsurl)
            ws.send(msg)
            resp=ws.recv()
            ws.close()
            return resp
    """
    @gen.coroutine
    def fct(ioloop,u,content):
        cnx=yield websocket_connect(u)
        cnx.write_message(msg)
        resp=yield cnx.read_message()
        cnx.close()
        ioloop.stop()
        ioloop.response=resp

    ioloop = IOLoop.instance()
    fct(ioloop,wsurl,msg)
    ioloop.start()
    return ioloop.response 

def _on_open(self):
    # Only proxy local connections.
    proxy_path = self.request.uri.replace(self._PATH_PREFIX, '/')
    proxy_url = '{}://{}{}'.format(_PROTOCOL_MAP[self.request.protocol],
                                   self.request.host, proxy_path)
    yield self._attach_auth_cookies()

    self.log.info('proxying WebSocket connection to: {}'.format(proxy_url))
    proxy_request = httpclient.HTTPRequest(
        url=proxy_url,
        method='GET',
        headers=self.request.headers,
        body=None,
        ca_certs=self.ca_certs)
    _modify_proxy_request_test_only(proxy_request)

    client = yield websocket.websocket_connect(
        proxy_request, on_message_callback=self._on_proxied_message)
    raise gen.Return(client) 

def test_proxied_post_base64_body(self):
    client = yield websocket.websocket_connect(
        self.get_ws_connection_request('http_over_websocket'))
    client.write_message(
        self.get_request_json(
            '/api/sessions',
            '1234',
            method='POST',
            body_base64=base64.b64encode(
                'the_data'.encode('utf-8')).decode('utf-8')))

    response_body = yield client.read_message()
    response = json.loads(response_body)
    self.assertEqual(200, response['status'])
    self.assertEqual('the_data',
                     base64.b64decode(response['data']).decode('utf-8'))
    self.assertEqual('1234', response['message_id'])
    self.assertTrue(response['done']) 

def test_websocket_close_buffered_data(self):
        ws = yield websocket_connect(
            'ws://127.0.0.1:%d/echo' % self.get_http_port())
        ws.write_message('hello')
        ws.write_message('world')
        # Close the underlying stream.
        ws.stream.close()
        yield self.close_future 

def test_missing_param(self):
    client = yield websocket.websocket_connect(
        self.get_ws_connection_request('http_over_websocket'))
    # Missing the 'method' parameter.
    client.write_message('{"path": "/api/sessions", "method": "GET"}')

    response_body = yield client.read_message()
    response = json.loads(response_body)
    self.assertEqual(400, response['status'])
    self.assertIn('body must contain', response['statusText'])
    self.assertTrue(response['done']) 

def test_invalid_protocol_version_requested(self):
    request = self.get_ws_connection_request('http_over_websocket')
    request.url += '?min_version=abc'

    client = yield websocket.websocket_connect(request)
    msg = yield client.read_message()
    # Message of None indicates that the connection has been closed.
    self.assertIsNone(msg)
    self.assertEqual(400, client.close_code)
    self.assertEqual('Invalid "min_version" provided: abc', client.close_reason) 

def test_websocket_callbacks(self):
        websocket_connect(
            'ws://127.0.0.1:%d/echo' % self.get_http_port(),
            io_loop=self.io_loop, callback=self.stop)
        ws = self.wait().result()
        ws.write_message('hello')
        ws.read_message(self.stop)
        response = self.wait().result()
        self.assertEqual(response, 'hello')
        self.close_future.add_done_callback(lambda f: self.stop())
        ws.close()
        self.wait() 

def test_websocket_network_fail(self):
        sock, port = bind_unused_port()
        sock.close()
        with self.assertRaises(IOError):
            with ExpectLog(gen_log, ".*"):
                yield websocket_connect(
                    'ws://127.0.0.1:%d/' % port,
                    io_loop=self.io_loop,
                    connect_timeout=3600) 

def test_diagnostic_handler_disallowed_cross_domain_origin(self):
    request = self.get_ws_connection_request('http_over_websocket/diagnose')
    request.headers.add('Origin', 'http://www.example.com')
    with self.assertRaises(httpclient.HTTPError) as e:
      yield websocket.websocket_connect(request)

    self.assertEqual(403, e.exception.code) 

def test_bad_json(self):
    client = yield websocket.websocket_connect(
        self.get_ws_connection_request('http_over_websocket'))
    client.write_message('invalid json')

    response_body = yield client.read_message()
    response = json.loads(response_body)
    self.assertEqual(400, response['status'])
    self.assertEqual('JSON input is required.', response['statusText'])
    self.assertTrue(response['done']) 

def test_max_chunk_size_64k(self):
    # Ensure that the size of the data propagated to the client is below 64K.
    # This size was chosen to ensure that Javascript JSON deserialization occurs
    # in a bounded amount of time.
    # Tornado itself does not make this parameter configurable, so add a
    # regression test here to ensure the default does not change. Reference:
    # https://github.com/tornadoweb/tornado/blob/master/tornado/http1connection.py#L76
    client = yield websocket.websocket_connect(
        self.get_ws_connection_request('http_over_websocket'))
    client.write_message(self.get_request_json('/api/largeresponse', '1234'))

    # Read all responses and ensure that they are smaller than the 64K limit.
    # HTTPS connections split the body into even smaller chunks.
    tornado_chunk_size = 65536
    responses = []
    while True:
      response_body = yield client.read_message()
      response = json.loads(response_body)
      responses.append(response)

      self.assertEqual('1234', response['message_id'])
      self.assertLessEqual(
          len(base64.b64decode(response['data'])), tornado_chunk_size)

      if response.get('done'):
        break 

def test_allowed_cross_domain_origin(self):
    request = self.get_ws_connection_request('http_over_websocket')
    request.headers.add('Origin', ALLOWED_ORIGIN)
    client = yield websocket.websocket_connect(request)
    self.assertIsNotNone(client) 

def test_disallowed_cross_domain_origin(self):
    request = self.get_ws_connection_request('http_over_websocket')
    request.headers.add('Origin', 'http://www.example.com')
    with self.assertRaises(httpclient.HTTPError) as e:
      yield websocket.websocket_connect(request)

    self.assertEqual(403, e.exception.code) 

def test_proxied_endpoint_has_error(self):
    client = yield websocket.websocket_connect(
        self.get_ws_connection_request('http_over_websocket'))
    # Force a 500 response. See _CANNED_RESPONSES
    client.write_message(
        self.get_request_json('/api/sessions?throw_500=1', '1234'))

    response_body = yield client.read_message()
    response = json.loads(response_body)
    self.assertEqual(500, response['status'])
    self.assertEqual('Server error', response['statusText'])
    self.assertEqual('1234', response['message_id'])
    self.assertTrue(response['done']) 

def test_request_for_invalid_url(self):
    client = yield websocket.websocket_connect(
        self.get_ws_connection_request('http_over_websocket'))
    client.write_message(self.get_request_json('/invalid/path', '1234'))

    response_body = yield client.read_message()
    response = json.loads(response_body)
    self.assertEqual(404, response['status'])
    self.assertEqual('1234', response['message_id'])
    self.assertTrue(response['done']) 

def test_streamed_response(self):
    client = yield websocket.websocket_connect(
        self.get_ws_connection_request('http_over_websocket'))
    client.write_message(self.get_request_json('/api/streamedresponse', '1234'))

    # The /api/streamedresponse handler yields a body split over two response.
    # The first will contain "first" as the body and contain metadata such as
    # status and headers. The second will contain "last" as the body and have
    # the done field set to true.
    first_response_body = yield client.read_message()
    first_response = json.loads(first_response_body)
    self.assertEqual(200, first_response['status'])
    self.assertEqual('1234', first_response['message_id'])
    self.assertEqual('first',
                     base64.b64decode(first_response['data']).decode('utf-8'))
    self.assertFalse(first_response['done'])

    second_response_body = yield client.read_message()
    second_response = json.loads(second_response_body)
    self.assertNotIn('status', second_response)
    self.assertNotIn('statusText', second_response)
    self.assertNotIn('headers', second_response)
    self.assertEqual('1234', second_response['message_id'])
    self.assertEqual('last',
                     base64.b64decode(second_response['data']).decode('utf-8'))
    self.assertTrue(second_response['done']) 

def test_proxied_post_no_body(self):
    client = yield websocket.websocket_connect(
        self.get_ws_connection_request('http_over_websocket'))
    client.write_message(
        self.get_request_json('/api/sessions', '1234', body=None))

    response_body = yield client.read_message()
    response = json.loads(response_body)
    self.assertEqual(200, response['status'])
    self.assertEqual('ok', base64.b64decode(response['data']).decode('utf-8'))
    self.assertEqual('1234', response['message_id'])
    self.assertTrue(response['done'])