1 /* 2 * Copyright (C) 2014 Square, Inc. 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 package com.squareup.okhttp.ws; 17 18 import com.squareup.okhttp.MediaType; 19 import com.squareup.okhttp.Response; 20 import com.squareup.okhttp.ResponseBody; 21 import com.squareup.okhttp.internal.ws.WebSocketReader; 22 import java.io.IOException; 23 import java.util.concurrent.BlockingQueue; 24 import java.util.concurrent.LinkedBlockingQueue; 25 import java.util.concurrent.TimeUnit; 26 import okio.Buffer; 27 28 import static com.squareup.okhttp.ws.WebSocket.BINARY; 29 import static com.squareup.okhttp.ws.WebSocket.TEXT; 30 import static org.junit.Assert.assertEquals; 31 import static org.junit.Assert.assertNotNull; 32 import static org.junit.Assert.assertTrue; 33 34 public final class WebSocketRecorder implements WebSocketReader.FrameCallback, WebSocketListener { 35 public interface MessageDelegate { 36 void onMessage(ResponseBody message) throws IOException; 37 } 38 39 private final BlockingQueue<Object> events = new LinkedBlockingQueue<>(); 40 private MessageDelegate delegate; 41 42 /** Sets a delegate for the next call to {@link #onMessage}. Cleared after invoked. */ 43 public void setNextMessageDelegate(MessageDelegate delegate) { 44 this.delegate = delegate; 45 } 46 47 @Override public void onOpen(WebSocket webSocket, Response response) { 48 } 49 50 @Override public void onMessage(ResponseBody message) throws IOException { 51 if (delegate != null) { 52 delegate.onMessage(message); 53 delegate = null; 54 } else { 55 Message event = new Message(message.contentType()); 56 message.source().readAll(event.buffer); 57 message.close(); 58 events.add(event); 59 } 60 } 61 62 @Override public void onPing(Buffer buffer) { 63 events.add(new Ping(buffer)); 64 } 65 66 @Override public void onPong(Buffer buffer) { 67 events.add(new Pong(buffer)); 68 } 69 70 @Override public void onClose(int code, String reason) { 71 events.add(new Close(code, reason)); 72 } 73 74 @Override public void onFailure(IOException e, Response response) { 75 events.add(e); 76 } 77 78 private Object nextEvent() { 79 try { 80 Object event = events.poll(10, TimeUnit.SECONDS); 81 if (event == null) { 82 throw new AssertionError("Timed out."); 83 } 84 return event; 85 } catch (InterruptedException e) { 86 throw new AssertionError(e); 87 } 88 } 89 90 public void assertTextMessage(String payload) throws IOException { 91 Message message = new Message(TEXT); 92 message.buffer.writeUtf8(payload); 93 Object actual = nextEvent(); 94 if (actual instanceof IOException) { 95 throw (IOException) actual; 96 } 97 assertEquals(message, actual); 98 } 99 100 public void assertBinaryMessage(byte[] payload) throws IOException { 101 Message message = new Message(BINARY); 102 message.buffer.write(payload); 103 Object actual = nextEvent(); 104 if (actual instanceof IOException) { 105 throw (IOException) actual; 106 } 107 assertEquals(message, actual); 108 } 109 110 public void assertPing(Buffer payload) throws IOException { 111 Object actual = nextEvent(); 112 if (actual instanceof IOException) { 113 throw (IOException) actual; 114 } 115 assertEquals(new Ping(payload), actual); 116 } 117 118 public void assertPong(Buffer payload) throws IOException { 119 Object actual = nextEvent(); 120 if (actual instanceof IOException) { 121 throw (IOException) actual; 122 } 123 assertEquals(new Pong(payload), actual); 124 } 125 126 public void assertClose(int code, String reason) throws IOException { 127 Object actual = nextEvent(); 128 if (actual instanceof IOException) { 129 throw (IOException) actual; 130 } 131 assertEquals(new Close(code, reason), actual); 132 } 133 134 public void assertFailure(Class<? extends IOException> cls, String message) { 135 Object event = nextEvent(); 136 String errorMessage = 137 "Expected [" + cls.getName() + ": " + message + "] but was [" + event + "]."; 138 assertNotNull(errorMessage, event); 139 assertEquals(errorMessage, cls, event.getClass()); 140 assertEquals(errorMessage, cls.cast(event).getMessage(), message); 141 } 142 143 public void assertExhausted() { 144 assertTrue("Remaining events: " + events, events.isEmpty()); 145 } 146 147 private static class Message { 148 public final MediaType mediaType; 149 public final Buffer buffer = new Buffer(); 150 151 private Message(MediaType mediaType) { 152 this.mediaType = mediaType; 153 } 154 155 @Override public String toString() { 156 return "Message[" + mediaType + " " + buffer + "]"; 157 } 158 159 @Override public int hashCode() { 160 return mediaType.hashCode() * 37 + buffer.hashCode(); 161 } 162 163 @Override public boolean equals(Object obj) { 164 if (obj instanceof Message) { 165 Message other = (Message) obj; 166 return mediaType.equals(other.mediaType) && buffer.equals(other.buffer); 167 } 168 return false; 169 } 170 } 171 172 private static class Ping { 173 public final Buffer buffer; 174 175 private Ping(Buffer buffer) { 176 this.buffer = buffer; 177 } 178 179 @Override public String toString() { 180 return "Ping[" + buffer + "]"; 181 } 182 183 @Override public int hashCode() { 184 return buffer.hashCode(); 185 } 186 187 @Override public boolean equals(Object obj) { 188 if (obj instanceof Ping) { 189 Ping other = (Ping) obj; 190 return buffer == null ? other.buffer == null : buffer.equals(other.buffer); 191 } 192 return false; 193 } 194 } 195 196 private static class Pong { 197 public final Buffer buffer; 198 199 private Pong(Buffer buffer) { 200 this.buffer = buffer; 201 } 202 203 @Override public String toString() { 204 return "Pong[" + buffer + "]"; 205 } 206 207 @Override public int hashCode() { 208 return buffer.hashCode(); 209 } 210 211 @Override public boolean equals(Object obj) { 212 if (obj instanceof Pong) { 213 Pong other = (Pong) obj; 214 return buffer == null ? other.buffer == null : buffer.equals(other.buffer); 215 } 216 return false; 217 } 218 } 219 220 private static class Close { 221 public final int code; 222 public final String reason; 223 224 private Close(int code, String reason) { 225 this.code = code; 226 this.reason = reason; 227 } 228 229 @Override public String toString() { 230 return "Close[" + code + " " + reason + "]"; 231 } 232 233 @Override public int hashCode() { 234 return code * 37 + reason.hashCode(); 235 } 236 237 @Override public boolean equals(Object obj) { 238 if (obj instanceof Close) { 239 Close other = (Close) obj; 240 return code == other.code && reason.equals(other.reason); 241 } 242 return false; 243 } 244 } 245 } 246
