ZkExternalResource.java

package com.netflix.suro.sink.kafka;

import org.I0Itec.zkclient.ZkClient;
import org.I0Itec.zkclient.exception.ZkMarshallingError;
import org.I0Itec.zkclient.serialize.ZkSerializer;
import org.apache.curator.test.TestingServer;
import org.junit.rules.ExternalResource;
import org.junit.rules.TemporaryFolder;

import java.io.IOException;
import java.io.UnsupportedEncodingException;

public class ZkExternalResource extends ExternalResource {
    private TestingServer zkServer;
    private ZkClient    zkClient;
    private TemporaryFolder tempDir = new TemporaryFolder();


    @Override
    protected void before() throws Throwable {
        tempDir.create();

        zkServer = new TestingServer();

        zkClient = new ZkClient("localhost:" + zkServer.getPort(), 20000, 20000, new ZkSerializer() {
            @Override
            public byte[] serialize(Object data) throws ZkMarshallingError {
                try {
                    return ((String)data).getBytes("UTF-8");
                } catch (UnsupportedEncodingException e) {
                    throw new RuntimeException(e);
                }
            }

            @Override
            public Object deserialize(byte[] bytes) throws ZkMarshallingError {
                if (bytes == null)
                    return null;
                try {
                    return new String(bytes, "UTF-8");
                } catch (UnsupportedEncodingException e) {
                    throw new RuntimeException(e);
                }
            }
        });
    }

    @Override
    protected void after() {
        if (zkServer != null) {
            try {
                zkServer.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        tempDir.delete();
    }

    public ZkClient getZkClient() {
        return zkClient;
    }
    public int getServerPort() { return zkServer.getPort(); }
    public String getConnectionString() { return zkServer.getConnectString(); }
}