We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Where assert statements are guarding against invalid ASGI messaging, use RuntimeError instead.
### Checklist

- [X] There are no similar issues or pull requests for this yet.
- [X] I discussed this idea on the [community chat](https://gitter.im/encode/community) and feedback is positive.

### Is your feature related to a problem? Please describe.

There are `assert` statements in the source code which raise a vague and hard to debug `AssertionError`. For example on [this line](https://github.com/encode/starlette/blob/f12c92a21500d484b3d48f965bb605c1bbe193bc/starlette/websockets.py#L58).

If some kind of exception (for example something along the lines of: `WebSocketMessageType`) were raised it would make debugging a lot clearer. I spent a lot more time than I should have just working out where exactly this `AssertionError` was coming from and what the root cause was.



### Describe the solution you would like.

This is by no means the right solution but at least it's an idea of the kind of thing that might help:

```python
class WebSocketMessageType(Exception):
    pass

class WebSocket(HTTPConnection):
    ...
    async def send(self, message: Message) -> None:
        """
        Send ASGI websocket messages, ensuring valid state transitions.
        """
        if self.application_state == WebSocketState.CONNECTING:
            message_type = message["type"]
            if message_type not in {"websocket.accept", "websocket.close"}:
                raise WebSocketMessageType("expected message_type to be websocket.accept or websocket.close")
            if message_type == "websocket.close":
                self.application_state = WebSocketState.DISCONNECTED
            else:
                self.application_state = WebSocketState.CONNECTED
            await self._send(message)
        elif self.application_state == WebSocketState.CONNECTED:
            message_type = message["type"]
            if message_type not in {"websocket.send", "websocket.close"}:
                raise WebSocketMessageType("expected message_type to be websocket.send or websocket.close")
            if message_type == "websocket.close":
                self.application_state = WebSocketState.DISCONNECTED
            await self._send(message)
        else:
            raise RuntimeError('Cannot call "send" once a close message has been sent.')
```

### Describe alternatives you considered

_No response_

### Additional context

The error I was seeing:

```
ERROR    root:a_file.py:31 {'message': 'Job processing failed', 'job': <Job coro=<<coroutine object a_class.a_method at 0x7f6d7a7c1ec0>>>, 'exception': AssertionError()}
NoneType: None
```

And this would be it with a `raise` statement: admittedly there is still no mention of `starlette` so a user would still have to diagnose that as the root cause.

```
ERROR    root:a_file.py:31 {'message': 'Job processing failed', 'job': <Job coro=<<coroutine object a_class.a_method at 0x7fb99c2ed940>>>, 'exception': WebSocketMessageType('expected message_type to be websocket.accept or websocket.close')}
NoneType: None
```

Also, I have no idea where that `NoneType: None` is coming from or what that means.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `starlette/websockets.py`
Content:
```
1 import enum
2 import json
3 import typing
4 
5 from starlette.requests import HTTPConnection
6 from starlette.types import Message, Receive, Scope, Send
7 
8 
9 class WebSocketState(enum.Enum):
10     CONNECTING = 0
11     CONNECTED = 1
12     DISCONNECTED = 2
13 
14 
15 class WebSocketDisconnect(Exception):
16     def __init__(self, code: int = 1000, reason: str = None) -> None:
17         self.code = code
18         self.reason = reason or ""
19 
20 
21 class WebSocket(HTTPConnection):
22     def __init__(self, scope: Scope, receive: Receive, send: Send) -> None:
23         super().__init__(scope)
24         assert scope["type"] == "websocket"
25         self._receive = receive
26         self._send = send
27         self.client_state = WebSocketState.CONNECTING
28         self.application_state = WebSocketState.CONNECTING
29 
30     async def receive(self) -> Message:
31         """
32         Receive ASGI websocket messages, ensuring valid state transitions.
33         """
34         if self.client_state == WebSocketState.CONNECTING:
35             message = await self._receive()
36             message_type = message["type"]
37             assert message_type == "websocket.connect"
38             self.client_state = WebSocketState.CONNECTED
39             return message
40         elif self.client_state == WebSocketState.CONNECTED:
41             message = await self._receive()
42             message_type = message["type"]
43             assert message_type in {"websocket.receive", "websocket.disconnect"}
44             if message_type == "websocket.disconnect":
45                 self.client_state = WebSocketState.DISCONNECTED
46             return message
47         else:
48             raise RuntimeError(
49                 'Cannot call "receive" once a disconnect message has been received.'
50             )
51 
52     async def send(self, message: Message) -> None:
53         """
54         Send ASGI websocket messages, ensuring valid state transitions.
55         """
56         if self.application_state == WebSocketState.CONNECTING:
57             message_type = message["type"]
58             assert message_type in {"websocket.accept", "websocket.close"}
59             if message_type == "websocket.close":
60                 self.application_state = WebSocketState.DISCONNECTED
61             else:
62                 self.application_state = WebSocketState.CONNECTED
63             await self._send(message)
64         elif self.application_state == WebSocketState.CONNECTED:
65             message_type = message["type"]
66             assert message_type in {"websocket.send", "websocket.close"}
67             if message_type == "websocket.close":
68                 self.application_state = WebSocketState.DISCONNECTED
69             await self._send(message)
70         else:
71             raise RuntimeError('Cannot call "send" once a close message has been sent.')
72 
73     async def accept(
74         self,
75         subprotocol: str = None,
76         headers: typing.Iterable[typing.Tuple[bytes, bytes]] = None,
77     ) -> None:
78         headers = headers or []
79 
80         if self.client_state == WebSocketState.CONNECTING:
81             # If we haven't yet seen the 'connect' message, then wait for it first.
82             await self.receive()
83         await self.send(
84             {"type": "websocket.accept", "subprotocol": subprotocol, "headers": headers}
85         )
86 
87     def _raise_on_disconnect(self, message: Message) -> None:
88         if message["type"] == "websocket.disconnect":
89             raise WebSocketDisconnect(message["code"])
90 
91     async def receive_text(self) -> str:
92         assert self.application_state == WebSocketState.CONNECTED
93         message = await self.receive()
94         self._raise_on_disconnect(message)
95         return message["text"]
96 
97     async def receive_bytes(self) -> bytes:
98         assert self.application_state == WebSocketState.CONNECTED
99         message = await self.receive()
100         self._raise_on_disconnect(message)
101         return message["bytes"]
102 
103     async def receive_json(self, mode: str = "text") -> typing.Any:
104         assert mode in ["text", "binary"]
105         assert self.application_state == WebSocketState.CONNECTED
106         message = await self.receive()
107         self._raise_on_disconnect(message)
108 
109         if mode == "text":
110             text = message["text"]
111         else:
112             text = message["bytes"].decode("utf-8")
113         return json.loads(text)
114 
115     async def iter_text(self) -> typing.AsyncIterator[str]:
116         try:
117             while True:
118                 yield await self.receive_text()
119         except WebSocketDisconnect:
120             pass
121 
122     async def iter_bytes(self) -> typing.AsyncIterator[bytes]:
123         try:
124             while True:
125                 yield await self.receive_bytes()
126         except WebSocketDisconnect:
127             pass
128 
129     async def iter_json(self) -> typing.AsyncIterator[typing.Any]:
130         try:
131             while True:
132                 yield await self.receive_json()
133         except WebSocketDisconnect:
134             pass
135 
136     async def send_text(self, data: str) -> None:
137         await self.send({"type": "websocket.send", "text": data})
138 
139     async def send_bytes(self, data: bytes) -> None:
140         await self.send({"type": "websocket.send", "bytes": data})
141 
142     async def send_json(self, data: typing.Any, mode: str = "text") -> None:
143         assert mode in ["text", "binary"]
144         text = json.dumps(data)
145         if mode == "text":
146             await self.send({"type": "websocket.send", "text": text})
147         else:
148             await self.send({"type": "websocket.send", "bytes": text.encode("utf-8")})
149 
150     async def close(self, code: int = 1000, reason: str = None) -> None:
151         await self.send(
152             {"type": "websocket.close", "code": code, "reason": reason or ""}
153         )
154 
155 
156 class WebSocketClose:
157     def __init__(self, code: int = 1000, reason: str = None) -> None:
158         self.code = code
159         self.reason = reason or ""
160 
161     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
162         await send(
163             {"type": "websocket.close", "code": self.code, "reason": self.reason}
164         )
165
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/starlette/websockets.py b/starlette/websockets.py
--- a/starlette/websockets.py
+++ b/starlette/websockets.py
@@ -34,13 +34,21 @@
         if self.client_state == WebSocketState.CONNECTING:
             message = await self._receive()
             message_type = message["type"]
-            assert message_type == "websocket.connect"
+            if message_type != "websocket.connect":
+                raise RuntimeError(
+                    'Expected ASGI message "websocket.connect", '
+                    f"but got {message_type!r}"
+                )
             self.client_state = WebSocketState.CONNECTED
             return message
         elif self.client_state == WebSocketState.CONNECTED:
             message = await self._receive()
             message_type = message["type"]
-            assert message_type in {"websocket.receive", "websocket.disconnect"}
+            if message_type not in {"websocket.receive", "websocket.disconnect"}:
+                raise RuntimeError(
+                    'Expected ASGI message "websocket.receive" or '
+                    f'"websocket.disconnect", but got {message_type!r}'
+                )
             if message_type == "websocket.disconnect":
                 self.client_state = WebSocketState.DISCONNECTED
             return message
@@ -55,7 +63,11 @@
         """
         if self.application_state == WebSocketState.CONNECTING:
             message_type = message["type"]
-            assert message_type in {"websocket.accept", "websocket.close"}
+            if message_type not in {"websocket.accept", "websocket.close"}:
+                raise RuntimeError(
+                    'Expected ASGI message "websocket.connect", '
+                    f"but got {message_type!r}"
+                )
             if message_type == "websocket.close":
                 self.application_state = WebSocketState.DISCONNECTED
             else:
@@ -63,7 +75,11 @@
             await self._send(message)
         elif self.application_state == WebSocketState.CONNECTED:
             message_type = message["type"]
-            assert message_type in {"websocket.send", "websocket.close"}
+            if message_type not in {"websocket.send", "websocket.close"}:
+                raise RuntimeError(
+                    'Expected ASGI message "websocket.send" or "websocket.close", '
+                    f"but got {message_type!r}"
+                )
             if message_type == "websocket.close":
                 self.application_state = WebSocketState.DISCONNECTED
             await self._send(message)
@@ -89,20 +105,30 @@
             raise WebSocketDisconnect(message["code"])
 
     async def receive_text(self) -> str:
-        assert self.application_state == WebSocketState.CONNECTED
+        if self.application_state != WebSocketState.CONNECTED:
+            raise RuntimeError(
+                'WebSocket is not connected. Need to call "accept" first.'
+            )
         message = await self.receive()
         self._raise_on_disconnect(message)
         return message["text"]
 
     async def receive_bytes(self) -> bytes:
-        assert self.application_state == WebSocketState.CONNECTED
+        if self.application_state != WebSocketState.CONNECTED:
+            raise RuntimeError(
+                'WebSocket is not connected. Need to call "accept" first.'
+            )
         message = await self.receive()
         self._raise_on_disconnect(message)
         return message["bytes"]
 
     async def receive_json(self, mode: str = "text") -> typing.Any:
-        assert mode in ["text", "binary"]
-        assert self.application_state == WebSocketState.CONNECTED
+        if mode not in {"text", "binary"}:
+            raise RuntimeError('The "mode" argument should be "text" or "binary".')
+        if self.application_state != WebSocketState.CONNECTED:
+            raise RuntimeError(
+                'WebSocket is not connected. Need to call "accept" first.'
+            )
         message = await self.receive()
         self._raise_on_disconnect(message)
 
@@ -140,7 +166,8 @@
         await self.send({"type": "websocket.send", "bytes": data})
 
     async def send_json(self, data: typing.Any, mode: str = "text") -> None:
-        assert mode in ["text", "binary"]
+        if mode not in {"text", "binary"}:
+            raise RuntimeError('The "mode" argument should be "text" or "binary".')
         text = json.dumps(data)
         if mode == "text":
             await self.send({"type": "websocket.send", "text": text})

{"golden_diff": "diff --git a/starlette/websockets.py b/starlette/websockets.py\n--- a/starlette/websockets.py\n+++ b/starlette/websockets.py\n@@ -34,13 +34,21 @@\n         if self.client_state == WebSocketState.CONNECTING:\n             message = await self._receive()\n             message_type = message[\"type\"]\n-            assert message_type == \"websocket.connect\"\n+            if message_type != \"websocket.connect\":\n+                raise RuntimeError(\n+                    'Expected ASGI message \"websocket.connect\", '\n+                    f\"but got {message_type!r}\"\n+                )\n             self.client_state = WebSocketState.CONNECTED\n             return message\n         elif self.client_state == WebSocketState.CONNECTED:\n             message = await self._receive()\n             message_type = message[\"type\"]\n-            assert message_type in {\"websocket.receive\", \"websocket.disconnect\"}\n+            if message_type not in {\"websocket.receive\", \"websocket.disconnect\"}:\n+                raise RuntimeError(\n+                    'Expected ASGI message \"websocket.receive\" or '\n+                    f'\"websocket.disconnect\", but got {message_type!r}'\n+                )\n             if message_type == \"websocket.disconnect\":\n                 self.client_state = WebSocketState.DISCONNECTED\n             return message\n@@ -55,7 +63,11 @@\n         \"\"\"\n         if self.application_state == WebSocketState.CONNECTING:\n             message_type = message[\"type\"]\n-            assert message_type in {\"websocket.accept\", \"websocket.close\"}\n+            if message_type not in {\"websocket.accept\", \"websocket.close\"}:\n+                raise RuntimeError(\n+                    'Expected ASGI message \"websocket.connect\", '\n+                    f\"but got {message_type!r}\"\n+                )\n             if message_type == \"websocket.close\":\n                 self.application_state = WebSocketState.DISCONNECTED\n             else:\n@@ -63,7 +75,11 @@\n             await self._send(message)\n         elif self.application_state == WebSocketState.CONNECTED:\n             message_type = message[\"type\"]\n-            assert message_type in {\"websocket.send\", \"websocket.close\"}\n+            if message_type not in {\"websocket.send\", \"websocket.close\"}:\n+                raise RuntimeError(\n+                    'Expected ASGI message \"websocket.send\" or \"websocket.close\", '\n+                    f\"but got {message_type!r}\"\n+                )\n             if message_type == \"websocket.close\":\n                 self.application_state = WebSocketState.DISCONNECTED\n             await self._send(message)\n@@ -89,20 +105,30 @@\n             raise WebSocketDisconnect(message[\"code\"])\n \n     async def receive_text(self) -> str:\n-        assert self.application_state == WebSocketState.CONNECTED\n+        if self.application_state != WebSocketState.CONNECTED:\n+            raise RuntimeError(\n+                'WebSocket is not connected. Need to call \"accept\" first.'\n+            )\n         message = await self.receive()\n         self._raise_on_disconnect(message)\n         return message[\"text\"]\n \n     async def receive_bytes(self) -> bytes:\n-        assert self.application_state == WebSocketState.CONNECTED\n+        if self.application_state != WebSocketState.CONNECTED:\n+            raise RuntimeError(\n+                'WebSocket is not connected. Need to call \"accept\" first.'\n+            )\n         message = await self.receive()\n         self._raise_on_disconnect(message)\n         return message[\"bytes\"]\n \n     async def receive_json(self, mode: str = \"text\") -> typing.Any:\n-        assert mode in [\"text\", \"binary\"]\n-        assert self.application_state == WebSocketState.CONNECTED\n+        if mode not in {\"text\", \"binary\"}:\n+            raise RuntimeError('The \"mode\" argument should be \"text\" or \"binary\".')\n+        if self.application_state != WebSocketState.CONNECTED:\n+            raise RuntimeError(\n+                'WebSocket is not connected. Need to call \"accept\" first.'\n+            )\n         message = await self.receive()\n         self._raise_on_disconnect(message)\n \n@@ -140,7 +166,8 @@\n         await self.send({\"type\": \"websocket.send\", \"bytes\": data})\n \n     async def send_json(self, data: typing.Any, mode: str = \"text\") -> None:\n-        assert mode in [\"text\", \"binary\"]\n+        if mode not in {\"text\", \"binary\"}:\n+            raise RuntimeError('The \"mode\" argument should be \"text\" or \"binary\".')\n         text = json.dumps(data)\n         if mode == \"text\":\n             await self.send({\"type\": \"websocket.send\", \"text\": text})\n", "issue": "Where assert statements are guarding against invalid ASGI messaging, use RuntimeError instead.\n### Checklist\r\n\r\n- [X] There are no similar issues or pull requests for this yet.\r\n- [X] I discussed this idea on the [community chat](https://gitter.im/encode/community) and feedback is positive.\r\n\r\n### Is your feature related to a problem? Please describe.\r\n\r\nThere are `assert` statements in the source code which raise a vague and hard to debug `AssertionError`. For example on [this line](https://github.com/encode/starlette/blob/f12c92a21500d484b3d48f965bb605c1bbe193bc/starlette/websockets.py#L58).\r\n\r\nIf some kind of exception (for example something along the lines of: `WebSocketMessageType`) were raised it would make debugging a lot clearer. I spent a lot more time than I should have just working out where exactly this `AssertionError` was coming from and what the root cause was.\r\n\r\n\r\n\r\n### Describe the solution you would like.\r\n\r\nThis is by no means the right solution but at least it's an idea of the kind of thing that might help:\r\n\r\n```python\r\nclass WebSocketMessageType(Exception):\r\n    pass\r\n\r\nclass WebSocket(HTTPConnection):\r\n    ...\r\n    async def send(self, message: Message) -> None:\r\n        \"\"\"\r\n        Send ASGI websocket messages, ensuring valid state transitions.\r\n        \"\"\"\r\n        if self.application_state == WebSocketState.CONNECTING:\r\n            message_type = message[\"type\"]\r\n            if message_type not in {\"websocket.accept\", \"websocket.close\"}:\r\n                raise WebSocketMessageType(\"expected message_type to be websocket.accept or websocket.close\")\r\n            if message_type == \"websocket.close\":\r\n                self.application_state = WebSocketState.DISCONNECTED\r\n            else:\r\n                self.application_state = WebSocketState.CONNECTED\r\n            await self._send(message)\r\n        elif self.application_state == WebSocketState.CONNECTED:\r\n            message_type = message[\"type\"]\r\n            if message_type not in {\"websocket.send\", \"websocket.close\"}:\r\n                raise WebSocketMessageType(\"expected message_type to be websocket.send or websocket.close\")\r\n            if message_type == \"websocket.close\":\r\n                self.application_state = WebSocketState.DISCONNECTED\r\n            await self._send(message)\r\n        else:\r\n            raise RuntimeError('Cannot call \"send\" once a close message has been sent.')\r\n```\r\n\r\n### Describe alternatives you considered\r\n\r\n_No response_\r\n\r\n### Additional context\r\n\r\nThe error I was seeing:\r\n\r\n```\r\nERROR    root:a_file.py:31 {'message': 'Job processing failed', 'job': <Job coro=<<coroutine object a_class.a_method at 0x7f6d7a7c1ec0>>>, 'exception': AssertionError()}\r\nNoneType: None\r\n```\r\n\r\nAnd this would be it with a `raise` statement: admittedly there is still no mention of `starlette` so a user would still have to diagnose that as the root cause.\r\n\r\n```\r\nERROR    root:a_file.py:31 {'message': 'Job processing failed', 'job': <Job coro=<<coroutine object a_class.a_method at 0x7fb99c2ed940>>>, 'exception': WebSocketMessageType('expected message_type to be websocket.accept or websocket.close')}\r\nNoneType: None\r\n```\r\n\r\nAlso, I have no idea where that `NoneType: None` is coming from or what that means.\n", "before_files": [{"content": "import enum\nimport json\nimport typing\n\nfrom starlette.requests import HTTPConnection\nfrom starlette.types import Message, Receive, Scope, Send\n\n\nclass WebSocketState(enum.Enum):\n    CONNECTING = 0\n    CONNECTED = 1\n    DISCONNECTED = 2\n\n\nclass WebSocketDisconnect(Exception):\n    def __init__(self, code: int = 1000, reason: str = None) -> None:\n        self.code = code\n        self.reason = reason or \"\"\n\n\nclass WebSocket(HTTPConnection):\n    def __init__(self, scope: Scope, receive: Receive, send: Send) -> None:\n        super().__init__(scope)\n        assert scope[\"type\"] == \"websocket\"\n        self._receive = receive\n        self._send = send\n        self.client_state = WebSocketState.CONNECTING\n        self.application_state = WebSocketState.CONNECTING\n\n    async def receive(self) -> Message:\n        \"\"\"\n        Receive ASGI websocket messages, ensuring valid state transitions.\n        \"\"\"\n        if self.client_state == WebSocketState.CONNECTING:\n            message = await self._receive()\n            message_type = message[\"type\"]\n            assert message_type == \"websocket.connect\"\n            self.client_state = WebSocketState.CONNECTED\n            return message\n        elif self.client_state == WebSocketState.CONNECTED:\n            message = await self._receive()\n            message_type = message[\"type\"]\n            assert message_type in {\"websocket.receive\", \"websocket.disconnect\"}\n            if message_type == \"websocket.disconnect\":\n                self.client_state = WebSocketState.DISCONNECTED\n            return message\n        else:\n            raise RuntimeError(\n                'Cannot call \"receive\" once a disconnect message has been received.'\n            )\n\n    async def send(self, message: Message) -> None:\n        \"\"\"\n        Send ASGI websocket messages, ensuring valid state transitions.\n        \"\"\"\n        if self.application_state == WebSocketState.CONNECTING:\n            message_type = message[\"type\"]\n            assert message_type in {\"websocket.accept\", \"websocket.close\"}\n            if message_type == \"websocket.close\":\n                self.application_state = WebSocketState.DISCONNECTED\n            else:\n                self.application_state = WebSocketState.CONNECTED\n            await self._send(message)\n        elif self.application_state == WebSocketState.CONNECTED:\n            message_type = message[\"type\"]\n            assert message_type in {\"websocket.send\", \"websocket.close\"}\n            if message_type == \"websocket.close\":\n                self.application_state = WebSocketState.DISCONNECTED\n            await self._send(message)\n        else:\n            raise RuntimeError('Cannot call \"send\" once a close message has been sent.')\n\n    async def accept(\n        self,\n        subprotocol: str = None,\n        headers: typing.Iterable[typing.Tuple[bytes, bytes]] = None,\n    ) -> None:\n        headers = headers or []\n\n        if self.client_state == WebSocketState.CONNECTING:\n            # If we haven't yet seen the 'connect' message, then wait for it first.\n            await self.receive()\n        await self.send(\n            {\"type\": \"websocket.accept\", \"subprotocol\": subprotocol, \"headers\": headers}\n        )\n\n    def _raise_on_disconnect(self, message: Message) -> None:\n        if message[\"type\"] == \"websocket.disconnect\":\n            raise WebSocketDisconnect(message[\"code\"])\n\n    async def receive_text(self) -> str:\n        assert self.application_state == WebSocketState.CONNECTED\n        message = await self.receive()\n        self._raise_on_disconnect(message)\n        return message[\"text\"]\n\n    async def receive_bytes(self) -> bytes:\n        assert self.application_state == WebSocketState.CONNECTED\n        message = await self.receive()\n        self._raise_on_disconnect(message)\n        return message[\"bytes\"]\n\n    async def receive_json(self, mode: str = \"text\") -> typing.Any:\n        assert mode in [\"text\", \"binary\"]\n        assert self.application_state == WebSocketState.CONNECTED\n        message = await self.receive()\n        self._raise_on_disconnect(message)\n\n        if mode == \"text\":\n            text = message[\"text\"]\n        else:\n            text = message[\"bytes\"].decode(\"utf-8\")\n        return json.loads(text)\n\n    async def iter_text(self) -> typing.AsyncIterator[str]:\n        try:\n            while True:\n                yield await self.receive_text()\n        except WebSocketDisconnect:\n            pass\n\n    async def iter_bytes(self) -> typing.AsyncIterator[bytes]:\n        try:\n            while True:\n                yield await self.receive_bytes()\n        except WebSocketDisconnect:\n            pass\n\n    async def iter_json(self) -> typing.AsyncIterator[typing.Any]:\n        try:\n            while True:\n                yield await self.receive_json()\n        except WebSocketDisconnect:\n            pass\n\n    async def send_text(self, data: str) -> None:\n        await self.send({\"type\": \"websocket.send\", \"text\": data})\n\n    async def send_bytes(self, data: bytes) -> None:\n        await self.send({\"type\": \"websocket.send\", \"bytes\": data})\n\n    async def send_json(self, data: typing.Any, mode: str = \"text\") -> None:\n        assert mode in [\"text\", \"binary\"]\n        text = json.dumps(data)\n        if mode == \"text\":\n            await self.send({\"type\": \"websocket.send\", \"text\": text})\n        else:\n            await self.send({\"type\": \"websocket.send\", \"bytes\": text.encode(\"utf-8\")})\n\n    async def close(self, code: int = 1000, reason: str = None) -> None:\n        await self.send(\n            {\"type\": \"websocket.close\", \"code\": code, \"reason\": reason or \"\"}\n        )\n\n\nclass WebSocketClose:\n    def __init__(self, code: int = 1000, reason: str = None) -> None:\n        self.code = code\n        self.reason = reason or \"\"\n\n    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:\n        await send(\n            {\"type\": \"websocket.close\", \"code\": self.code, \"reason\": self.reason}\n        )\n", "path": "starlette/websockets.py"}], "after_files": [{"content": "import enum\nimport json\nimport typing\n\nfrom starlette.requests import HTTPConnection\nfrom starlette.types import Message, Receive, Scope, Send\n\n\nclass WebSocketState(enum.Enum):\n    CONNECTING = 0\n    CONNECTED = 1\n    DISCONNECTED = 2\n\n\nclass WebSocketDisconnect(Exception):\n    def __init__(self, code: int = 1000, reason: str = None) -> None:\n        self.code = code\n        self.reason = reason or \"\"\n\n\nclass WebSocket(HTTPConnection):\n    def __init__(self, scope: Scope, receive: Receive, send: Send) -> None:\n        super().__init__(scope)\n        assert scope[\"type\"] == \"websocket\"\n        self._receive = receive\n        self._send = send\n        self.client_state = WebSocketState.CONNECTING\n        self.application_state = WebSocketState.CONNECTING\n\n    async def receive(self) -> Message:\n        \"\"\"\n        Receive ASGI websocket messages, ensuring valid state transitions.\n        \"\"\"\n        if self.client_state == WebSocketState.CONNECTING:\n            message = await self._receive()\n            message_type = message[\"type\"]\n            if message_type != \"websocket.connect\":\n                raise RuntimeError(\n                    'Expected ASGI message \"websocket.connect\", '\n                    f\"but got {message_type!r}\"\n                )\n            self.client_state = WebSocketState.CONNECTED\n            return message\n        elif self.client_state == WebSocketState.CONNECTED:\n            message = await self._receive()\n            message_type = message[\"type\"]\n            if message_type not in {\"websocket.receive\", \"websocket.disconnect\"}:\n                raise RuntimeError(\n                    'Expected ASGI message \"websocket.receive\" or '\n                    f'\"websocket.disconnect\", but got {message_type!r}'\n                )\n            if message_type == \"websocket.disconnect\":\n                self.client_state = WebSocketState.DISCONNECTED\n            return message\n        else:\n            raise RuntimeError(\n                'Cannot call \"receive\" once a disconnect message has been received.'\n            )\n\n    async def send(self, message: Message) -> None:\n        \"\"\"\n        Send ASGI websocket messages, ensuring valid state transitions.\n        \"\"\"\n        if self.application_state == WebSocketState.CONNECTING:\n            message_type = message[\"type\"]\n            if message_type not in {\"websocket.accept\", \"websocket.close\"}:\n                raise RuntimeError(\n                    'Expected ASGI message \"websocket.connect\", '\n                    f\"but got {message_type!r}\"\n                )\n            if message_type == \"websocket.close\":\n                self.application_state = WebSocketState.DISCONNECTED\n            else:\n                self.application_state = WebSocketState.CONNECTED\n            await self._send(message)\n        elif self.application_state == WebSocketState.CONNECTED:\n            message_type = message[\"type\"]\n            if message_type not in {\"websocket.send\", \"websocket.close\"}:\n                raise RuntimeError(\n                    'Expected ASGI message \"websocket.send\" or \"websocket.close\", '\n                    f\"but got {message_type!r}\"\n                )\n            if message_type == \"websocket.close\":\n                self.application_state = WebSocketState.DISCONNECTED\n            await self._send(message)\n        else:\n            raise RuntimeError('Cannot call \"send\" once a close message has been sent.')\n\n    async def accept(\n        self,\n        subprotocol: str = None,\n        headers: typing.Iterable[typing.Tuple[bytes, bytes]] = None,\n    ) -> None:\n        headers = headers or []\n\n        if self.client_state == WebSocketState.CONNECTING:\n            # If we haven't yet seen the 'connect' message, then wait for it first.\n            await self.receive()\n        await self.send(\n            {\"type\": \"websocket.accept\", \"subprotocol\": subprotocol, \"headers\": headers}\n        )\n\n    def _raise_on_disconnect(self, message: Message) -> None:\n        if message[\"type\"] == \"websocket.disconnect\":\n            raise WebSocketDisconnect(message[\"code\"])\n\n    async def receive_text(self) -> str:\n        if self.application_state != WebSocketState.CONNECTED:\n            raise RuntimeError(\n                'WebSocket is not connected. Need to call \"accept\" first.'\n            )\n        message = await self.receive()\n        self._raise_on_disconnect(message)\n        return message[\"text\"]\n\n    async def receive_bytes(self) -> bytes:\n        if self.application_state != WebSocketState.CONNECTED:\n            raise RuntimeError(\n                'WebSocket is not connected. Need to call \"accept\" first.'\n            )\n        message = await self.receive()\n        self._raise_on_disconnect(message)\n        return message[\"bytes\"]\n\n    async def receive_json(self, mode: str = \"text\") -> typing.Any:\n        if mode not in {\"text\", \"binary\"}:\n            raise RuntimeError('The \"mode\" argument should be \"text\" or \"binary\".')\n        if self.application_state != WebSocketState.CONNECTED:\n            raise RuntimeError(\n                'WebSocket is not connected. Need to call \"accept\" first.'\n            )\n        message = await self.receive()\n        self._raise_on_disconnect(message)\n\n        if mode == \"text\":\n            text = message[\"text\"]\n        else:\n            text = message[\"bytes\"].decode(\"utf-8\")\n        return json.loads(text)\n\n    async def iter_text(self) -> typing.AsyncIterator[str]:\n        try:\n            while True:\n                yield await self.receive_text()\n        except WebSocketDisconnect:\n            pass\n\n    async def iter_bytes(self) -> typing.AsyncIterator[bytes]:\n        try:\n            while True:\n                yield await self.receive_bytes()\n        except WebSocketDisconnect:\n            pass\n\n    async def iter_json(self) -> typing.AsyncIterator[typing.Any]:\n        try:\n            while True:\n                yield await self.receive_json()\n        except WebSocketDisconnect:\n            pass\n\n    async def send_text(self, data: str) -> None:\n        await self.send({\"type\": \"websocket.send\", \"text\": data})\n\n    async def send_bytes(self, data: bytes) -> None:\n        await self.send({\"type\": \"websocket.send\", \"bytes\": data})\n\n    async def send_json(self, data: typing.Any, mode: str = \"text\") -> None:\n        if mode not in {\"text\", \"binary\"}:\n            raise RuntimeError('The \"mode\" argument should be \"text\" or \"binary\".')\n        text = json.dumps(data)\n        if mode == \"text\":\n            await self.send({\"type\": \"websocket.send\", \"text\": text})\n        else:\n            await self.send({\"type\": \"websocket.send\", \"bytes\": text.encode(\"utf-8\")})\n\n    async def close(self, code: int = 1000, reason: str = None) -> None:\n        await self.send(\n            {\"type\": \"websocket.close\", \"code\": code, \"reason\": reason or \"\"}\n        )\n\n\nclass WebSocketClose:\n    def __init__(self, code: int = 1000, reason: str = None) -> None:\n        self.code = code\n        self.reason = reason or \"\"\n\n    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:\n        await send(\n            {\"type\": \"websocket.close\", \"code\": self.code, \"reason\": self.reason}\n        )\n", "path": "starlette/websockets.py"}]}