package com.proofpoint.http.client;

import com.google.common.base.Charsets;
import com.google.common.base.Throwables;
import com.google.common.io.CharStreams;
import com.proofpoint.json.JsonCodec;

import java.io.IOException;
import java.io.InputStreamReader;
import java.net.ConnectException;

public class JsonResponseHandler<T> implements ResponseHandler<T, RuntimeException>
{
    private static final MediaType MEDIA_TYPE_JSON = MediaType.create("application", "json");

    public static <T> JsonResponseHandler<T> createJsonResponseHandler(JsonCodec<T> jsonCodec)
    {
        return new JsonResponseHandler<T>(jsonCodec);
    }

    private final JsonCodec<T> jsonCodec;

    private JsonResponseHandler(JsonCodec<T> jsonCodec)
    {
        this.jsonCodec = jsonCodec;
    }

    @Override
    public RuntimeException handleException(Request request, Exception exception)
    {
        if (exception instanceof ConnectException) {
            return new RuntimeException("Server refused connection: " + request.getUri().toASCIIString());
        }
        return Throwables.propagate(exception);
    }

    @Override
    public T handle(Request request, Response response)
    {
        if (response.getStatusCode() / 100 != 2) {
            throw new UnexpectedResponseException(request, response);
        }
        String contentType = response.getHeader("Content-Type");
        if (!MediaType.parse(contentType).is(MEDIA_TYPE_JSON)) {
            throw new UnexpectedResponseException("Expected application/json response from server but got " + contentType, request, response);
        }
        try {
            String json = CharStreams.toString(new InputStreamReader(response.getInputStream(), Charsets.UTF_8));
            T value = jsonCodec.fromJson(json);
            return value;
        }
        catch (IOException e) {
            throw new RuntimeException("Error reading response from server");
        }
    }

}

@@ -141,21 +141,110 @@
 mon.
-io.CharStream
+collect.ImmutableSet;%0Aimport com.google.common.io.CharStreams;%0Aimport com.google.common.primitives.Int
 s;%0Ai
@@ -376,16 +376,38 @@
 ception;
+%0Aimport java.util.Set;
 %0A%0Apublic
@@ -752,142 +752,851 @@
    p
-rivate final JsonCodec%3CT%3E jsonCodec;%0A%0A    private JsonResponseHandler(JsonCodec%3CT%3E jsonCodec)%0A    %7B%0A        this.jsonCodec = jsonCodec
+ublic static %3CT%3E JsonResponseHandler%3CT%3E createJsonResponseHandler(JsonCodec%3CT%3E jsonCodec, int firstSuccessfulResponseCode, int... otherSuccessfulResponseCodes)%0A    %7B%0A        return new JsonResponseHandler%3CT%3E(jsonCodec, firstSuccessfulResponseCode, otherSuccessfulResponseCodes);%0A    %7D%0A%0A    private final JsonCodec%3CT%3E jsonCodec;%0A    private final Set%3CInteger%3E successfulResponseCodes;%0A%0A    private JsonResponseHandler(JsonCodec%3CT%3E jsonCodec)%0A    %7B%0A        this(jsonCodec, 200, 201, 202, 203, 204, 205, 206);%0A    %7D%0A%0A    private JsonResponseHandler(JsonCodec%3CT%3E jsonCodec, int firstSuccessfulResponseCode, int... otherSuccessfulResponseCodes)%0A    %7B%0A        this.jsonCodec = jsonCodec;%0A        this.successfulResponseCodes = ImmutableSet.%3CInteger%3Ebuilder().add(firstSuccessfulResponseCode).addAll(Ints.asList(otherSuccessfulResponseCodes)).build()
 ;%0A  
@@ -2014,16 +2014,50 @@
      if 
+(!successfulResponseCodes.contains
 (respons
@@ -2077,19 +2077,9 @@
 de()
- / 100 != 2
+)
 ) %7B%0A
@@ -2128,24 +2128,231 @@
 ception(
-request,
+%0A                    String.format(%22Expected response code to be %25s, but was %25d: %25s%22, successfulResponseCodes, response.getStatusCode(), response.getStatusMessage()),%0A                    request,%0A                   
  respons
@@ -2986,11 +2986,10 @@
 %7D%0A    %7D%0A
-%0A
 %7D%0A

Make list of successful response codes configurable