package com.facebook.thrift.direct_server.tests;

import static org.junit.Assert.*;

import com.facebook.thrift.direct_server.TDirectServer;
import com.facebook.thrift.protocol.TBinaryProtocol;
import com.facebook.thrift.transport.TFramedTransport;
import com.facebook.thrift.transport.TSocket;
import com.facebook.thrift.transport.TTransportException;
import org.junit.BeforeClass;
import org.junit.Test;

import com.facebook.fbcode.fb303.fb_status;

public class JUnitTestHsHaServer {

  private static final String HOST = "localhost";
  private static final int PORT = 19191;
  private static final int NUM_THREADS = 1;
  private static final int MAX_PENDING = 1;

  private static JavaSimpleService.Client newClient() throws Exception {
    TFramedTransport trans = new TFramedTransport(new TSocket(HOST, PORT));
    trans.open();
    return new JavaSimpleService.Client(new TBinaryProtocol(trans));
  }

  @BeforeClass
  public static void setUpBeforeClass() throws Exception {
    JavaSimpleService.Processor proc =
        new JavaSimpleService.Processor(new SimpleServiceHandler());
    final TDirectServer server =
        TDirectServer.asHsHaServer(PORT, NUM_THREADS, MAX_PENDING, proc);
    new Thread() {
      @Override
      public void run() {
        server.serve();
      }
    }.start();
    final int NUM_ITERS = 20, MS_PER_ITER = 100;
    for (int i = 0; i < NUM_ITERS; i++) {
      try {
        if (newClient().getStatus() == fb_status.ALIVE) {
          return;
        }
      } catch (TTransportException e) {
        if (i == NUM_ITERS - 1) {
          throw e;
        }
      }
      Thread.sleep(MS_PER_ITER);
    }
  }

  @Test
  public void simple() throws Exception {
    assertEquals(fb_status.ALIVE, newClient().getStatus());
    assertEquals("", newClient().getString(0));
    assertEquals("xxxxx", newClient().getString(5));
  }

  @Test
  public void largeResponse() throws Exception {
    final int LENGTH = 1024 * 1024;
    assertEquals(LENGTH, newClient().getString(LENGTH).length());
  }
}

@@ -346,16 +346,60 @@
 eption;%0A
+import java.util.concurrent.CountDownLatch;%0A
 import o
@@ -596,22 +596,16 @@
 tic 
-final 
 int PORT
  = 1
@@ -604,16 +604,8 @@
 PORT
- = 19191
 ;%0A  
@@ -932,16 +932,504 @@
 );%0A  %7D%0A%0A
+  private static class PortGetter implements Runnable %7B%0A    private CountDownLatch latch_;%0A    private TDirectServer server_;%0A    private int port_;%0A%0A    public PortGetter(CountDownLatch latch) %7B%0A      latch_ = latch;%0A    %7D%0A%0A    public void setServer(TDirectServer server) %7B%0A      server_ = server;%0A    %7D%0A%0A    public int getPort() %7B%0A      return port_;%0A    %7D%0A%0A    @Override%0A    public void run() %7B%0A      port_ = server_.directServer().getLocalPort();%0A      latch_.countDown();%0A    %7D%0A  %7D%0A%0A
   @Befor
@@ -1598,24 +1598,125 @@
 Handler());%0A
+    CountDownLatch latch = new CountDownLatch(1);%0A    PortGetter portGetter = new PortGetter(latch);%0A
     final TD
@@ -1771,20 +1771,17 @@
 aServer(
-PORT
+0
 , NUM_TH
@@ -1804,16 +1804,62 @@
 NG, proc
+, portGetter);%0A    portGetter.setServer(server
 );%0A    n
@@ -1957,24 +1957,76 @@
  %7D.start();%0A
+    latch.await();%0A    PORT = portGetter.getPort();%0A
     final in

Fix DirectServer test