diff --git a/axum/src/extract/ws.rs b/axum/src/extract/ws.rs
index 8965bd73..482e6916 100644
--- a/axum/src/extract/ws.rs
+++ b/axum/src/extract/ws.rs
@@ -40,28 +40,28 @@
 //!
 //! ```
 //! use axum::{
-//!     extract::ws::{WebSocketUpgrade, WebSocket},
+//!     extract::{ws::{WebSocketUpgrade, WebSocket}, State},
 //!     response::Response,
 //!     routing::get,
-//!     Extension, Router,
+//!     Router,
 //! };
 //!
 //! #[derive(Clone)]
-//! struct State {
+//! struct AppState {
 //!     // ...
 //! }
 //!
-//! async fn handler(ws: WebSocketUpgrade, Extension(state): Extension<State>) -> Response {
+//! async fn handler(ws: WebSocketUpgrade, State(state): State<AppState>) -> Response {
 //!     ws.on_upgrade(|socket| handle_socket(socket, state))
 //! }
 //!
-//! async fn handle_socket(socket: WebSocket, state: State) {
+//! async fn handle_socket(socket: WebSocket, state: AppState) {
 //!     // ...
 //! }
 //!
 //! let app = Router::new()
 //!     .route("/ws", get(handler))
-//!     .layer(Extension(State { /* ... */ }));
+//!     .with_state(AppState { /* ... */ });
 //! # async {
 //! # axum::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap();
 //! # };