net.codestory.rest.RestAssert Maven / Gradle / Ivy

Go to download
/**
 * Copyright (C) 2013-2015 [email protected]
 *
 * 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 net.codestory.rest;

import com.squareup.okhttp.*;

import java.io.IOException;
import java.net.Proxy;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.UnaryOperator;

import static java.util.function.UnaryOperator.identity;

public class RestAssert {
  private final String url;
  private final UnaryOperator configureClient;
  private final UnaryOperator configureRequest;

  RestAssert(String url) {
    this(url, identity(), identity());
  }

  private RestAssert(String url, UnaryOperator configureClient, UnaryOperator configureRequest) {
    this.url = url;
    this.configureRequest = configureRequest;
    this.configureClient = configureClient;
  }

  // Creation
  RestAssert withRequest(UnaryOperator configure) {
    return new RestAssert(url, configureClient, request -> configure.apply(configureRequest.apply(request)));
  }

  private RestAssert withClient(UnaryOperator configure) {
    return new RestAssert(url, client -> configure.apply(configureClient.apply(client)), configureRequest);
  }

  // Extraction

  public Response response() {
    try {
      RestResponse call = call();
      return new Response(call.contentType(), call.bodyAsString(), call.code());
    } catch (IOException e) {
      throw new RuntimeException("Unable to query: " + url, e);
    }
  }

  // Configuration
  public RestAssert withHeader(String name, String value) {
    return withRequest(addHeader(name, value));
  }

  public RestAssert withPreemptiveAuthentication(String login, String password) {
    return withRequest(addBasicAuthHeader(login, password));
  }

  public RestAssert withAuthentication(String login, String password) {
    return withClient(setAuthenticator(new Authenticator() {
      AtomicInteger tries = new AtomicInteger(0);

      @Override
      public Request authenticate(Proxy proxy, com.squareup.okhttp.Response response) {
        if (tries.getAndIncrement() > 0) {
          return null;
        }
        return addBasicAuthHeader(login, password).apply(response.request().newBuilder()).build();
      }

      @Override
      public Request authenticateProxy(Proxy proxy, com.squareup.okhttp.Response response) {
        return null;
      }
    }));
  }

  // Assertions
  public Should should() {
    try {
      return new RestResponseShould(call(), false);
    } catch (IOException e) {
      throw new RuntimeException("Unable to query: " + url, e);
    }
  }

  // Client modification
  private static UnaryOperator setAuthenticator(Authenticator authenticator) {
    return client -> client.setAuthenticator(authenticator);
  }

  // Call
  private RestResponse call() throws IOException {
    return RestResponse.call(url, configureClient, configureRequest);
  }

  // Request configuration
  private static UnaryOperator addBasicAuthHeader(String login, String password) {
    return addHeader("Authorization", Credentials.basic(login, password));
  }

  private static UnaryOperator addHeader(String name, String value) {
    return request -> request.addHeader(name, value);
  }
}