diff --git a/crates/api_clients/grok_client/src/requests/listen_websocket/messages/websocket_server_message.rs b/crates/api_clients/grok_client/src/requests/listen_websocket/messages/websocket_server_message.rs
index bb5289431d..20b38ef2db 100644
--- a/crates/api_clients/grok_client/src/requests/listen_websocket/messages/websocket_server_message.rs
+++ b/crates/api_clients/grok_client/src/requests/listen_websocket/messages/websocket_server_message.rs
@@ -89,6 +89,14 @@ mod tests {
     let data = std::fs::read_to_string(path)?;
     let message: WebsocketServerMessage = serde_json::from_str(&data)?;
 
+    match message {
+      WebsocketServerMessage::ImageData(image_data) => {
+        assert_eq!(image_data.percentage_complete, Some(50));
+        assert!(image_data.url.is_some());
+      },
+      _ => panic!("Expected ImageData message"),
+    }
+
     Ok(())
   }
 
@@ -98,6 +106,13 @@ mod tests {
     let data = std::fs::read_to_string(path)?;
     let message: WebsocketServerMessage = serde_json::from_str(&data)?;
 
+    match message {
+      WebsocketServerMessage::JsonData(image_data) => {
+        assert_eq!(image_data.percentage_complete, Some(100));
+      },
+      _ => panic!("Expected JsonData message"),
+    }
+
     Ok(())
   }
 
@@ -107,6 +122,13 @@ mod tests {
     let data = std::fs::read_to_string(path)?;
     let message: WebsocketServerMessage = serde_json::from_str(&data)?;
 
+    match message {
+      WebsocketServerMessage::JsonData(image_data) => {
+        assert_eq!(image_data.percentage_complete, Some(0));
+      },
+      _ => panic!("Expected JsonData message"),
+    }
+
     Ok(())
   }