...27import org.openqa.selenium.remote.http.HttpClient;28import org.openqa.selenium.remote.http.HttpRequest;29import org.openqa.selenium.remote.http.HttpResponse;30import org.openqa.selenium.remote.http.Message;31import org.openqa.selenium.remote.http.TextMessage;32import org.openqa.selenium.remote.http.WebSocket;33import org.openqa.selenium.support.ui.FluentWait;34import org.openqa.selenium.testing.Safely;35import java.util.LinkedList;36import java.util.List;37import java.util.Optional;38import java.util.concurrent.CountDownLatch;39import java.util.concurrent.Executors;40import java.util.concurrent.ScheduledExecutorService;41import java.util.concurrent.atomic.AtomicBoolean;42import java.util.function.BiFunction;43import java.util.function.Consumer;44import static java.util.concurrent.TimeUnit.MILLISECONDS;45import static java.util.concurrent.TimeUnit.SECONDS;46import static org.assertj.core.api.Assertions.assertThat;47import static org.openqa.selenium.remote.http.Contents.utf8String;48import static org.openqa.selenium.remote.http.HttpMethod.GET;49public class WebSocketServingTest {50 private Server<?> server;51 @After52 public void shutDown() {53 Safely.safelyCall(() -> server.stop());54 }55 @Test(expected = ConnectionFailedException.class)56 public void clientShouldThrowAnExceptionIfUnableToConnectToAWebSocketEndPoint() {57 server = new NettyServer(defaultOptions(), req -> new HttpResponse()).start();58 HttpClient client = HttpClient.Factory.createDefault().createClient(server.getUrl());59 client.openSocket(new HttpRequest(GET, "/does-not-exist"), new WebSocket.Listener());60 }61 @Test62 public void shouldUseUriToChooseWhichWebSocketHandlerToUse() throws InterruptedException {63 AtomicBoolean foo = new AtomicBoolean(false);64 AtomicBoolean bar = new AtomicBoolean(false);65 BiFunction<String, Consumer<Message>, Optional<Consumer<Message>>> factory = (str, sink) -> {66 if ("/foo".equals(str)) {67 return Optional.of(msg -> {68 foo.set(true);69 sink.accept(new TextMessage("Foo called"));70 });71 } else {72 return Optional.of(msg -> {73 bar.set(true);74 sink.accept(new TextMessage("Bar called"));75 });76 }77 };78 server = new NettyServer(79 defaultOptions(),80 req -> new HttpResponse(),81 factory82 ).start();83 CountDownLatch latch = new CountDownLatch(1);84 HttpClient client = HttpClient.Factory.createDefault().createClient(server.getUrl());85 WebSocket fooSocket = client.openSocket(new HttpRequest(GET, "/foo"), new WebSocket.Listener() {86 @Override87 public void onText(CharSequence data) {88 System.out.println("Called!");89 latch.countDown();90 }91 });92 fooSocket.sendText("Hello, World!");93 latch.await(2, SECONDS);94 assertThat(foo.get()).isTrue();95 assertThat(bar.get()).isFalse();96 }97 @Test98 public void shouldStillBeAbleToServeHttpTraffic() {99 server = new NettyServer(100 defaultOptions(),101 req -> new HttpResponse().setContent(utf8String("Brie!")),102 (uri, sink) -> {103 if ("/foo".equals(uri)) {104 return Optional.of(msg -> sink.accept(new TextMessage("Peas!")));105 }106 return Optional.empty();107 }).start();108 HttpClient client = HttpClient.Factory.createDefault().createClient(server.getUrl());109 HttpResponse res = client.execute(new HttpRequest(GET, "/cheese"));110 assertThat(Contents.string(res)).isEqualTo("Brie!");111 }112 @Test113 public void shouldPropagateCloseMessage() throws InterruptedException {114 CountDownLatch latch = new CountDownLatch(1);115 server = new NettyServer(116 defaultOptions(),117 req -> new HttpResponse(),118 (uri, sink) -> Optional.of(socket -> latch.countDown())).start();119 HttpClient client = HttpClient.Factory.createDefault().createClient(server.getUrl());120 WebSocket socket = client.openSocket(new HttpRequest(GET, "/cheese"), new WebSocket.Listener());121 socket.close();122 latch.await(2, SECONDS);123 }124 @Test125 public void webSocketHandlersShouldBeAbleToFireMoreThanOneMessage() {126 server = new NettyServer(127 defaultOptions(),128 req -> new HttpResponse(),129 (uri, sink) -> Optional.of(msg -> {130 sink.accept(new TextMessage("beyaz peynir"));131 sink.accept(new TextMessage("cheddar"));132 })).start();133 HttpClient client = HttpClient.Factory.createDefault().createClient(server.getUrl());134 List<String> messages = new LinkedList<>();135 WebSocket socket = client.openSocket(new HttpRequest(GET, "/cheese"), new WebSocket.Listener() {136 @Override137 public void onText(CharSequence data) {138 messages.add(data.toString());139 }140 });141 socket.send(new TextMessage("Hello"));142 new FluentWait<>(messages).until(msgs -> msgs.size() == 2);143 }144 public void serverShouldBeAbleToPushAMessageWithoutNeedingTheClientToSendAMessage() throws InterruptedException {145 class MyHandler implements Consumer<Message> {146 private final Consumer<Message> sink;147 private final ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();148 public MyHandler(Consumer<Message> sink) {149 this.sink = sink;150 // Send a message every 250ms151 executor.scheduleAtFixedRate(152 () -> sink.accept(new TextMessage("Calling home.")),153 100,154 250,155 MILLISECONDS);156 }157 @Override158 public void accept(Message message) {159 // Do nothing160 }161 }162 server = new NettyServer(163 defaultOptions(),164 req -> new HttpResponse(),165 (uri, sink) -> Optional.of(new MyHandler(sink))).start();166 CountDownLatch latch = new CountDownLatch(2);...