zipkin.junit.ZipkinRule Maven / Gradle / Ivy

Go to download
/**
 * Copyright 2015-2016 The OpenZipkin Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 * in compliance with the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License
 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 * or implied. See the License for the specific language governing permissions and limitations under
 * the License.
 */
package zipkin.junit;

import java.io.IOException;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import okhttp3.mockwebserver.Dispatcher;
import okhttp3.mockwebserver.MockResponse;
import okhttp3.mockwebserver.MockWebServer;
import okhttp3.mockwebserver.RecordedRequest;
import org.junit.rules.TestRule;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;
import zipkin.InMemorySpanStore;
import zipkin.Span;

import static okhttp3.mockwebserver.SocketPolicy.KEEP_OPEN;

/**
 * Starts up a local Zipkin server, listening for http requests on {@link #httpUrl}.
 *
 * This can be used to test instrumentation. For example, you can POST spans directly to this
 * server.
 *
 * See http://openzipkin.github.io/zipkin-api/#/
 */
public final class ZipkinRule implements TestRule {

  private final InMemorySpanStore store = new InMemorySpanStore();
  private final MockWebServer server = new MockWebServer();
  private final BlockingQueue failureQueue = new LinkedBlockingQueue<>();

  public ZipkinRule() {
    Dispatcher dispatcher = new Dispatcher() {
      final ZipkinDispatcher successDispatch = new ZipkinDispatcher(store, server);

      @Override
      public MockResponse dispatch(RecordedRequest request) throws InterruptedException {
        MockResponse maybeFailure = failureQueue.poll();
        if (maybeFailure != null) return maybeFailure;
        return successDispatch.dispatch(request);
      }

      @Override
      public MockResponse peek() {
        MockResponse maybeFailure = failureQueue.peek();
        if (maybeFailure != null) return maybeFailure;
        return new MockResponse().setSocketPolicy(KEEP_OPEN);
      }
    };
    server.setDispatcher(dispatcher);
  }

  /** Use this to connect. The zipkin v1 interface will be under "/api/v1" */
  public String httpUrl() {
    return String.format("http://%s:%s", server.getHostName(), server.getPort());
  }

  /** Use this to see how many requests you've sent to any zipkin http endpoint. */
  public int httpRequestCount() {
    return server.getRequestCount();
  }

  /**
   * Stores the given spans directly, to setup preconditions for a test.
   *
   * 
For example, if you are testing what happens when instrumentation adds a child to a trace,
   * you'd add the parent here.
   */
  public ZipkinRule storeSpans(Iterable spans) {
    store.accept(spans.iterator());
    return this;
  }

  /**
   * Adds a one-time failure to the http endpoint.
   *
   * 
Ex. If you want to test that you don't repeatedly send bad data, you could send a 400 back.
   *
   * 
{@code
   * zipkin.enqueueFailure(sendErrorResponse(400, "bad format"));
   * }
   *
   * @param failure type of failure the next call to the http endpoint responds with
   */
  public ZipkinRule enqueueFailure(HttpFailure failure) {
    failureQueue.add(failure.response);
    return this;
  }

  /** Retrieves all traces this zipkin server has received. */
  public List> getTraces() {
    return store.getTracesByIds(store.traceIds());
  }

  /**
   * Used to manually start the server.
   *
   * @param httpPort choose 0 to select an available port
   */
  public void start(int httpPort) throws IOException {
    server.start(httpPort);
  }

  /** Used to manually stop the server. */
  public void shutdown() throws IOException {
    server.shutdown();
  }

  @Override
  public Statement apply(Statement base, Description description) {
    return server.apply(base, description);
  }
}