1/*
2 * Copyright (C) 2015 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 */
16package com.squareup.okhttp.ws;
17
18import com.squareup.okhttp.OkHttpClient;
19import com.squareup.okhttp.Request;
20import com.squareup.okhttp.Response;
21import com.squareup.okhttp.internal.Version;
22import java.io.IOException;
23import java.util.concurrent.CountDownLatch;
24import java.util.concurrent.ExecutorService;
25import java.util.concurrent.Executors;
26import java.util.concurrent.TimeUnit;
27import java.util.concurrent.atomic.AtomicLong;
28import java.util.concurrent.atomic.AtomicReference;
29import okio.Buffer;
30import okio.BufferedSource;
31
32/**
33 * Exercises the web socket implementation against the
34 * <a href="http://autobahn.ws/testsuite/">Autobahn Testsuite</a>.
35 */
36public final class AutobahnTester {
37  private static final String HOST = "ws://localhost:9001";
38
39  public static void main(String... args) throws IOException {
40    new AutobahnTester().run();
41  }
42
43  final OkHttpClient client = new OkHttpClient();
44
45  private WebSocketCall newWebSocket(String path) {
46    Request request = new Request.Builder().url(HOST + path).build();
47    return WebSocketCall.create(client, request);
48  }
49
50  public void run() throws IOException {
51    try {
52      long count = getTestCount();
53      System.out.println("Test count: " + count);
54
55      for (long number = 1; number <= count; number++) {
56        runTest(number, count);
57      }
58
59      updateReports();
60    } finally {
61      client.getDispatcher().getExecutorService().shutdown();
62    }
63  }
64
65  private void runTest(final long number, final long count) throws IOException {
66    final CountDownLatch latch = new CountDownLatch(1);
67    newWebSocket("/runCase?case=" + number + "&agent=" + Version.userAgent()) //
68        .enqueue(new WebSocketListener() {
69          private final ExecutorService sendExecutor = Executors.newSingleThreadExecutor();
70          private WebSocket webSocket;
71
72          @Override public void onOpen(WebSocket webSocket, Response response) {
73            System.out.println("Executing test case " + number + "/" + count);
74            this.webSocket = webSocket;
75          }
76
77          @Override public void onMessage(BufferedSource payload, final WebSocket.PayloadType type)
78              throws IOException {
79            final Buffer buffer = new Buffer();
80            payload.readAll(buffer);
81            payload.close();
82
83            sendExecutor.execute(new Runnable() {
84              @Override public void run() {
85                try {
86                  webSocket.sendMessage(type, buffer);
87                } catch (IOException e) {
88                  e.printStackTrace();
89                }
90              }
91            });
92          }
93
94          @Override public void onPong(Buffer payload) {
95          }
96
97          @Override public void onClose(int code, String reason) {
98            sendExecutor.shutdown();
99            latch.countDown();
100          }
101
102          @Override public void onFailure(IOException e, Response response) {
103            latch.countDown();
104          }
105        });
106    try {
107      if (!latch.await(10, TimeUnit.SECONDS)) {
108        throw new IllegalStateException("Timed out waiting for count.");
109      }
110    } catch (InterruptedException e) {
111      throw new AssertionError();
112    }
113  }
114
115  private long getTestCount() throws IOException {
116    final CountDownLatch latch = new CountDownLatch(1);
117    final AtomicLong countRef = new AtomicLong();
118    final AtomicReference<IOException> failureRef = new AtomicReference<>();
119    newWebSocket("/getCaseCount").enqueue(new WebSocketListener() {
120      @Override public void onOpen(WebSocket webSocket, Response response) {
121      }
122
123      @Override public void onMessage(BufferedSource payload, WebSocket.PayloadType type)
124          throws IOException {
125        countRef.set(payload.readDecimalLong());
126        payload.close();
127      }
128
129      @Override public void onPong(Buffer payload) {
130      }
131
132      @Override public void onClose(int code, String reason) {
133        latch.countDown();
134      }
135
136      @Override public void onFailure(IOException e, Response response) {
137        failureRef.set(e);
138        latch.countDown();
139      }
140    });
141    try {
142      if (!latch.await(10, TimeUnit.SECONDS)) {
143        throw new IllegalStateException("Timed out waiting for count.");
144      }
145    } catch (InterruptedException e) {
146      throw new AssertionError();
147    }
148    IOException failure = failureRef.get();
149    if (failure != null) {
150      throw failure;
151    }
152    return countRef.get();
153  }
154
155  private void updateReports() {
156    final CountDownLatch latch = new CountDownLatch(1);
157    newWebSocket("/updateReports?agent=" + Version.userAgent()).enqueue(new WebSocketListener() {
158      @Override public void onOpen(WebSocket webSocket, Response response) {
159      }
160
161      @Override public void onMessage(BufferedSource payload, WebSocket.PayloadType type)
162          throws IOException {
163      }
164
165      @Override public void onPong(Buffer payload) {
166      }
167
168      @Override public void onClose(int code, String reason) {
169        latch.countDown();
170      }
171
172      @Override public void onFailure(IOException e, Response response) {
173        latch.countDown();
174      }
175    });
176    try {
177      if (!latch.await(10, TimeUnit.SECONDS)) {
178        throw new IllegalStateException("Timed out waiting for count.");
179      }
180    } catch (InterruptedException e) {
181      throw new AssertionError();
182    }
183  }
184}
185