Make redirects work with the async API.

okhttp/src/main/java/com/squareup/okhttp/Job.java

@@ -101,25 +101,23 @@ final class Job implements Runnable {
 
       engine.readResponse();
 
-      Response engineResponse = engine.getResponse();
-      Response response = engineResponse.newBuilder()
-          .body(new Dispatcher.RealResponseBody(engineResponse, engine.getResponseBody()))
-          .redirectedBy(redirectedBy)
-          .build();
-
+      Response response = engine.getResponse();
       Request redirect = processResponse(engine, response);
 
       if (redirect == null) {
         engine.automaticallyReleaseConnectionToPool();
-        return response;
+        return response.newBuilder()
+            .body(new Dispatcher.RealResponseBody(response, engine.getResponseBody()))
+            .redirectedBy(redirectedBy)
+            .build();
       }
 
       // TODO: fail if too many redirects
       // TODO: fail if not following redirects
-      // TODO: release engine
+      engine.release(false);
 
       connection = sameConnection(request, redirect) ? engine.getConnection() : null;
-      redirectedBy = response;
+      redirectedBy = response.newBuilder().redirectedBy(redirectedBy).build(); // Chained.
       request = redirect;
     }
   }

okhttp/src/test/java/com/squareup/okhttp/AsyncApiTest.java

@@ -71,6 +71,26 @@ public final class AsyncApiTest {
     assertTrue(server.takeRequest().getHeaders().contains("User-Agent: AsyncApiTest"));
   }
 
+  @Test public void connectionPooling() throws Exception {
+    server.enqueue(new MockResponse().setBody("abc"));
+    server.enqueue(new MockResponse().setBody("def"));
+    server.enqueue(new MockResponse().setBody("ghi"));
+    server.play();
+
+    client.enqueue(new Request.Builder().url(server.getUrl("/a")).build(), receiver);
+    receiver.await(server.getUrl("/a")).assertBody("abc");
+
+    client.enqueue(new Request.Builder().url(server.getUrl("/b")).build(), receiver);
+    receiver.await(server.getUrl("/b")).assertBody("def");
+
+    client.enqueue(new Request.Builder().url(server.getUrl("/c")).build(), receiver);
+    receiver.await(server.getUrl("/c")).assertBody("ghi");
+
+    assertEquals(0, server.takeRequest().getSequenceNumber());
+    assertEquals(1, server.takeRequest().getSequenceNumber());
+    assertEquals(2, server.takeRequest().getSequenceNumber());
+  }
+
   @Test public void tls() throws Exception {
     server.useHttps(sslContext.getSocketFactory(), false);
     server.enqueue(new MockResponse()
@@ -130,4 +150,32 @@ public final class AsyncApiTest {
     receiver.await(request2.url()).assertCode(200).assertBody("A");
     assertEquals("v1", server.takeRequest().getHeader("If-None-Match"));
   }
+
+  @Test public void redirect() throws Exception {
+    server.enqueue(new MockResponse()
+        .setResponseCode(301)
+        .addHeader("Location: /b")
+        .addHeader("Test", "Redirect from /a to /b")
+        .setBody("/a has moved!"));
+    server.enqueue(new MockResponse()
+        .setResponseCode(302)
+        .addHeader("Location: /c")
+        .addHeader("Test", "Redirect from /b to /c")
+        .setBody("/b has moved!"));
+    server.enqueue(new MockResponse().setBody("C"));
+    server.play();
+
+    Request request = new Request.Builder().url(server.getUrl("/a")).build();
+    client.enqueue(request, receiver);
+
+    receiver.await(server.getUrl("/c"))
+        .assertCode(200)
+        .assertBody("C")
+        .redirectedBy()
+        .assertCode(302)
+        .assertContainsHeaders("Test: Redirect from /b to /c")
+        .redirectedBy()
+        .assertCode(301)
+        .assertContainsHeaders("Test: Redirect from /a to /b");
+  }
 }

okhttp/src/test/java/com/squareup/okhttp/RecordedResponse.java

@@ -70,4 +70,15 @@ public class RecordedResponse {
     assertEquals(0, handshake.localCertificates().size());
     return this;
   }
+
+  /**
+   * Asserts that the current response was redirected and returns a new recorded
+   * response for the original request.
+   */
+  public RecordedResponse redirectedBy() {
+    Response redirectedBy = response.redirectedBy();
+    assertNotNull(redirectedBy);
+    assertNull(redirectedBy.body());
+    return new RecordedResponse(redirectedBy.request(), redirectedBy, null, null);
+  }
 }