text: import asyncio
  import base64
  import hashlib
  import os
  import unittest
  from unittest import mock
  
  import aiohttp
  from aiohttp import errors, hdrs, websocket, websocket_client
  
  
  class TestWebSocketClient(unittest.TestCase):
  
      def setUp(self):
          self.loop = asyncio.new_event_loop()
          asyncio.set_event_loop(None)
  
          self.key_data = os.urandom(16)
          self.key = base64.b64encode(self.key_data)
          self.ws_key = base64.b64encode(
              hashlib.sha1(self.key + websocket.WS_KEY).digest()).decode()
  
      def tearDown(self):
          self.loop.close()
  
      @mock.patch('aiohttp.client.os')
      @mock.patch('aiohttp.client.ClientSession.request')
      def test_ws_connect(self, m_req, m_os):
          resp = mock.Mock()
          resp.status = 101
          resp.headers = {
              hdrs.UPGRADE: hdrs.WEBSOCKET,
              hdrs.CONNECTION: hdrs.UPGRADE,
              hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key,
              hdrs.SEC_WEBSOCKET_PROTOCOL: 'chat'
          }
          m_os.urandom.return_value = self.key_data
          m_req.return_value = asyncio.Future(loop=self.loop)
          m_req.return_value.set_result(resp)
  
          res = self.loop.run_until_complete(
              aiohttp.ws_connect(
                  'http://test.org',
                  protocols=('t1', 't2', 'chat'),
                  loop=self.loop))
  
          self.assertIsInstance(res, websocket_client.ClientWebSocketResponse)
          self.assertEqual(res.protocol, 'chat')
          self.assertNotIn(hdrs.ORIGIN, m_req.call_args[1]["headers"])
  
      @mock.patch('aiohttp.client.os')
      @mock.patch('aiohttp.client.ClientSession.request')
      def test_ws_connect_with_origin(self, m_req, m_os):
          resp = mock.Mock()
          resp.status = 403
          m_os.urandom.return_value = self.key_data
          m_req.return_value = asyncio.Future(loop=self.loop)
          m_req.return_value.set_result(resp)
  
          origin = 'https://example.org/page.html'
          with self.assertRaises(errors.WSServerHandshakeError):
              self.loop.run_until_complete(
                  aiohttp.ws_connect(
                      'http://test.org',
                      loop=self.loop,
                      origin=origin))
  
          self.assertIn(hdrs.ORIGIN, m_req.call_args[1]["headers"])
          self.assertEqual(m_req.call_args[1]["headers"][hdrs.ORIGIN], origin)
  
      @mock.patch('aiohttp.client.os')
      @mock.patch('aiohttp.client.ClientSession.request')
      def test_ws_connect_custom_response(self, m_req, m_os):
  
          class CustomResponse(websocket_client.ClientWebSocketResponse):
              def read(self, decode=False):
                  return 'customized!'
  
          resp = mock.Mock()
          resp.status = 101
          resp.headers = {
              hdrs.UPGRADE: hdrs.WEBSOCKET,
              hdrs.CONNECTION: hdrs.UPGRADE,
              hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key,
          }
          m_os.urandom.return_value = self.key_data
          m_req.return_value = asyncio.Future(loop=self.loop)
          m_req.return_value.set_result(resp)
  
          res = self.loop.run_until_complete(
              aiohttp.ws_connect(
                  'http://test.org',
                  ws_response_class=CustomResponse,
                  loop=self.loop))
  
          self.assertEqual(res.read(), 'customized!')
  
      @mock.patch('aiohttp.client.os')
      @mock.patch('aiohttp.client.ClientSession.request')
      def test_ws_connect_global_loop(self, m_req, m_os):
          asyncio.set_event_loop(self.loop)
  
          resp = mock.Mock()
          resp.status = 101
          resp.headers = {
              hdrs.UPGRADE: hdrs.WEBSOCKET,
              hdrs.CONNECTION: hdrs.UPGRADE,
              hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key
          }
          m_os.urandom.return_value = self.key_data
          m_req.return_value = asyncio.Future(loop=self.loop)
          m_req.return_value.set_result(resp)
  
          resp = self.loop.run_until_complete(
              aiohttp.ws_connect('http://test.org'))
          self.assertIs(resp._loop, self.loop)
  
          asyncio.set_event_loop(None)
  
      @mock.patch('aiohttp.client.os')
      @mock.patch('aiohttp.client.ClientSession.request')
      def test_ws_connect_err_status(self, m_req, m_os):
          resp = mock.Mock()
          resp.status = 500
          resp.headers = {
              hdrs.UPGRADE: hdrs.WEBSOCKET,
              hdrs.CONNECTION: hdrs.UPGRADE,
              hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key
          }
          m_os.urandom.return_value = self.key_data
          m_req.return_value = asyncio.Future(loop=self.loop)
          m_req.return_value.set_result(resp)
  
          with self.assertRaises(errors.WSServerHandshakeError) as ctx:
              self.loop.run_until_complete(
                  aiohttp.ws_connect(
                      'http://test.org',
                      protocols=('t1', 't2', 'chat'),
                      loop=self.loop))
          self.assertEqual(
              ctx.exception.message, 'Invalid response status')
  
      @mock.patch('aiohttp.client.os')
      @mock.patch('aiohttp.client.ClientSession.request')
      def test_ws_connect_err_upgrade(self, m_req, m_os):
          resp = mock.Mock()
          resp.status = 101
          resp.headers = {
              hdrs.UPGRADE: 'test',
              hdrs.CONNECTION: hdrs.UPGRADE,
              hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key
          }
          m_os.urandom.return_value = self.key_data
          m_req.return_value = asyncio.Future(loop=self.loop)
          m_req.return_value.set_result(resp)
  
          with self.assertRaises(errors.WSServerHandshakeError) as ctx:
              self.loop.run_until_complete(
                  aiohttp.ws_connect(
                      'http://test.org',
                      protocols=('t1', 't2', 'chat'),
                      loop=self.loop))
          self.assertEqual(
              ctx.exception.message, 'Invalid upgrade header')
  
      @mock.patch('aiohttp.client.os')
      @mock.patch('aiohttp.client.ClientSession.request')
      def test_ws_connect_err_conn(self, m_req, m_os):
          resp = mock.Mock()
          resp.status = 101
          resp.headers = {
              hdrs.UPGRADE: hdrs.WEBSOCKET,
              hdrs.CONNECTION: 'close',
              hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key
          }
          m_os.urandom.return_value = self.key_data
          m_req.return_value = asyncio.Future(loop=self.loop)
          m_req.return_value.set_result(resp)
  
          with self.assertRaises(errors.WSServerHandshakeError) as ctx:
              self.loop.run_until_complete(
                  aiohttp.ws_connect(
                      'http://test.org',
                      protocols=('t1', 't2', 'chat'),
                      loop=self.loop))
          self.assertEqual(
              ctx.exception.message, 'Invalid connection header')
  
      @mock.patch('aiohttp.client.os')
      @mock.patch('aiohttp.client.ClientSession.request')
      def test_ws_connect_err_challenge(self, m_req, m_os):
          resp = mock.Mock()
          resp.status = 101
          resp.headers = {
              hdrs.UPGRADE: hdrs.WEBSOCKET,
              hdrs.CONNECTION: hdrs.UPGRADE,
              hdrs.SEC_WEBSOCKET_ACCEPT: 'asdfasdfasdfasdfasdfasdf'
          }
          m_os.urandom.return_value = self.key_data
          m_req.return_value = asyncio.Future(loop=self.loop)
          m_req.return_value.set_result(resp)
  
          with self.assertRaises(errors.WSServerHandshakeError) as ctx:
              self.loop.run_until_complete(
                  aiohttp.ws_connect(
                      'http://test.org',
                      protocols=('t1', 't2', 'chat'),
                      loop=self.loop))
          self.assertEqual(
              ctx.exception.message, 'Invalid challenge response')
  
      @mock.patch('aiohttp.client.WebSocketWriter')
      @mock.patch('aiohttp.client.os')
      @mock.patch('aiohttp.client.ClientSession.request')
      def test_close(self, m_req, m_os, WebSocketWriter):
          resp = mock.Mock()
          resp.status = 101
          resp.headers = {
              hdrs.UPGRADE: hdrs.WEBSOCKET,
              hdrs.CONNECTION: hdrs.UPGRADE,
              hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key,
          }
          m_os.urandom.return_value = self.key_data
          m_req.return_value = asyncio.Future(loop=self.loop)
          m_req.return_value.set_result(resp)
          writer = WebSocketWriter.return_value = mock.Mock()
          reader = resp.connection.reader.set_parser.return_value = mock.Mock()
  
          resp = self.loop.run_until_complete(
              aiohttp.ws_connect('http://test.org', loop=self.loop))
          self.assertFalse(resp.closed)
  
          msg = websocket.Message(websocket.MSG_CLOSE, b'', b'')
          reader.read.return_value = asyncio.Future(loop=self.loop)
          reader.read.return_value.set_result(msg)
  
          res = self.loop.run_until_complete(resp.close())
          writer.close.assert_called_with(1000, b'')
          self.assertTrue(resp.closed)
          self.assertTrue(res)
          self.assertIsNone(resp.exception())
  
          # idempotent
          res = self.loop.run_until_complete(resp.close())
          self.assertFalse(res)
          self.assertEqual(writer.close.call_count, 1)
  
      @mock.patch('aiohttp.client.WebSocketWriter')
      @mock.patch('aiohttp.client.os')
      @mock.patch('aiohttp.client.ClientSession.request')
      def test_close_exc(self, m_req, m_os, WebSocketWriter):
          resp = mock.Mock()
          resp.status = 101
          resp.headers = {
              hdrs.UPGRADE: hdrs.WEBSOCKET,
              hdrs.CONNECTION: hdrs.UPGRADE,
              hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key,
          }
          m_os.urandom.return_value = self.key_data
          m_req.return_value = asyncio.Future(loop=self.loop)
          m_req.return_value.set_result(resp)
          WebSocketWriter.return_value = mock.Mock()
          reader = resp.connection.reader.set_parser.return_value = mock.Mock()
  
          resp = self.loop.run_until_complete(
              aiohttp.ws_connect(
                  'http://test.org', loop=self.loop))
          self.assertFalse(resp.closed)
  
          exc = ValueError()
          reader.read.return_value = asyncio.Future(loop=self.loop)
          reader.read.return_value.set_exception(exc)
  
          self.loop.run_until_complete(resp.close())
          self.assertTrue(resp.closed)
          self.assertIs(resp.exception(), exc)
  
      @mock.patch('aiohttp.client.WebSocketWriter')
      @mock.patch('aiohttp.client.os')
      @mock.patch('aiohttp.client.ClientSession.request')
      def test_close_exc2(self, m_req, m_os, WebSocketWriter):
          resp = mock.Mock()
          resp.status = 101
          resp.headers = {
              hdrs.UPGRADE: hdrs.WEBSOCKET,
              hdrs.CONNECTION: hdrs.UPGRADE,
              hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key,
          }
          m_os.urandom.return_value = self.key_data
          m_req.return_value = asyncio.Future(loop=self.loop)
          m_req.return_value.set_result(resp)
          writer = WebSocketWriter.return_value = mock.Mock()
          resp.connection.reader.set_parser.return_value = mock.Mock()
  
          resp = self.loop.run_until_complete(
              aiohttp.ws_connect(
                  'http://test.org', loop=self.loop))
          self.assertFalse(resp.closed)
  
          exc = ValueError()
          writer.close.side_effect = exc
  
          self.loop.run_until_complete(resp.close())
          self.assertTrue(resp.closed)
          self.assertIs(resp.exception(), exc)
  
          resp._closed = False
          writer.close.side_effect = asyncio.CancelledError()
          self.assertRaises(asyncio.CancelledError,
                            self.loop.run_until_complete, resp.close())
  
      @mock.patch('aiohttp.client.WebSocketWriter')
      @mock.patch('aiohttp.client.os')
      @mock.patch('aiohttp.client.ClientSession.request')
      def test_send_data_after_close(self, m_req, m_os, WebSocketWriter):
          resp = mock.Mock()
          resp.status = 101
          resp.headers = {
              hdrs.UPGRADE: hdrs.WEBSOCKET,
              hdrs.CONNECTION: hdrs.UPGRADE,
              hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key,
          }
          m_os.urandom.return_value = self.key_data
          m_req.return_value = asyncio.Future(loop=self.loop)
          m_req.return_value.set_result(resp)
          WebSocketWriter.return_value = mock.Mock()
  
          resp = self.loop.run_until_complete(
              aiohttp.ws_connect(
                  'http://test.org', loop=self.loop))
          resp._closed = True
  
          self.assertRaises(RuntimeError, resp.ping)
          self.assertRaises(RuntimeError, resp.pong)
          self.assertRaises(RuntimeError, resp.send_str, 's')
          self.assertRaises(RuntimeError, resp.send_bytes, b'b')
  
      @mock.patch('aiohttp.client.WebSocketWriter')
      @mock.patch('aiohttp.client.os')
      @mock.patch('aiohttp.client.ClientSession.request')
      def test_send_data_type_errors(self, m_req, m_os, WebSocketWriter):
          resp = mock.Mock()
          resp.status = 101
          resp.headers = {
              hdrs.UPGRADE: hdrs.WEBSOCKET,
              hdrs.CONNECTION: hdrs.UPGRADE,
              hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key,
          }
          m_os.urandom.return_value = self.key_data
          m_req.return_value = asyncio.Future(loop=self.loop)
          m_req.return_value.set_result(resp)
          WebSocketWriter.return_value = mock.Mock()
  
          resp = self.loop.run_until_complete(
              aiohttp.ws_connect(
                  'http://test.org', loop=self.loop))
  
          self.assertRaises(TypeError, resp.send_str, b's')
          self.assertRaises(TypeError, resp.send_bytes, 'b')
  
      @mock.patch('aiohttp.client.WebSocketWriter')
      @mock.patch('aiohttp.client.os')
      @mock.patch('aiohttp.client.ClientSession.request')
      def test_reader_read_exception(self, m_req, m_os, WebSocketWriter):
          hresp = mock.Mock()
          hresp.status = 101
          hresp.headers = {
              hdrs.UPGRADE: hdrs.WEBSOCKET,
              hdrs.CONNECTION: hdrs.UPGRADE,
              hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key,
          }
          m_os.urandom.return_value = self.key_data
          m_req.return_value = asyncio.Future(loop=self.loop)
          m_req.return_value.set_result(hresp)
          WebSocketWriter.return_value = mock.Mock()
          reader = hresp.connection.reader.set_parser.return_value = mock.Mock()
  
          resp = self.loop.run_until_complete(
              aiohttp.ws_connect(
                  'http://test.org', loop=self.loop))
  
          exc = ValueError()
          reader.read.return_value = asyncio.Future(loop=self.loop)
          reader.read.return_value.set_exception(exc)
  
          msg = self.loop.run_until_complete(resp.receive())
          self.assertEqual(msg.tp, aiohttp.MsgType.error)
          self.assertIs(resp.exception(), exc)
  
      def test_receive_runtime_err(self):
          resp = websocket_client.ClientWebSocketResponse(
              mock.Mock(), mock.Mock(), mock.Mock(), mock.Mock(), 10.0,
              True, True, self.loop)
          resp._waiting = True
  
          self.assertRaises(
              RuntimeError, self.loop.run_until_complete, resp.receive())
  