AbstractDockerClientTest.java

package com.github.dockerjava.client;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.StringWriter;
import java.lang.reflect.Method;
import java.net.DatagramSocket;
import java.net.ServerSocket;
import java.util.ArrayList;
import java.util.List;

import org.apache.commons.io.IOUtils;
import org.apache.commons.io.LineIterator;
import org.apache.commons.lang.StringUtils;
import org.hamcrest.FeatureMatcher;
import org.hamcrest.Matcher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.Assert;
import org.testng.ITestResult;

import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.command.InspectContainerResponse;
import com.github.dockerjava.api.command.InspectContainerResponse.Mount;
import com.github.dockerjava.api.exception.DockerException;
import com.github.dockerjava.api.exception.NotFoundException;
import com.github.dockerjava.api.model.Frame;
import com.github.dockerjava.api.model.Network;
import com.github.dockerjava.api.model.Volume;
import com.github.dockerjava.core.DockerClientBuilder;
import com.github.dockerjava.core.DefaultDockerClientConfig;
import com.github.dockerjava.core.TestDockerCmdExecFactory;
import com.github.dockerjava.core.command.BuildImageResultCallback;
import com.github.dockerjava.core.command.LogContainerResultCallback;
import com.github.dockerjava.core.command.PullImageResultCallback;

public abstract class AbstractDockerClientTest extends Assert {

    public static final Logger LOG = LoggerFactory.getLogger(AbstractDockerClientTest.class);

    protected DockerClient dockerClient;

    protected TestDockerCmdExecFactory dockerCmdExecFactory = initTestDockerCmdExecFactory();

    protected TestDockerCmdExecFactory initTestDockerCmdExecFactory() {
        return new TestDockerCmdExecFactory(
                DockerClientBuilder.getDefaultDockerCmdExecFactory());
    }

    public void beforeTest() throws Exception {

        LOG.info("======================= BEFORETEST =======================");
        LOG.info("Connecting to Docker server");
        dockerClient = DockerClientBuilder.getInstance(config())
                .withDockerCmdExecFactory(dockerCmdExecFactory)
                .build();

        try {
            dockerClient.inspectImageCmd("busybox").exec();
        } catch (NotFoundException e) {
            LOG.info("Pulling image 'busybox'");
            // need to block until image is pulled completely
            dockerClient.pullImageCmd("busybox").withTag("latest").exec(new PullImageResultCallback()).awaitSuccess();
        }

        assertNotNull(dockerClient);
        LOG.info("======================= END OF BEFORETEST =======================\n\n");
    }

    private DefaultDockerClientConfig config() {
        return config(null);
    }

    protected DefaultDockerClientConfig config(String password) {
        DefaultDockerClientConfig.Builder builder = DefaultDockerClientConfig.createDefaultConfigBuilder()
                .withRegistryUrl("https://index.docker.io/v1/");
        if (password != null) {
            builder = builder.withRegistryPassword(password);
        }

        return builder.build();
    }

    public void afterTest() {
        LOG.info("======================= END OF AFTERTEST =======================");
    }

    public void beforeMethod(Method method) {
        LOG.info(String.format("################################## STARTING %s ##################################",
                method.getName()));
    }

    public void afterMethod(ITestResult result) {

        for (String container : dockerCmdExecFactory.getContainerNames()) {
            LOG.info("Cleaning up temporary container {}", container);

            try {
                dockerClient.removeContainerCmd(container).withForce(true).exec();
            } catch (DockerException ignore) {
                // ignore.printStackTrace();
            }
        }

        for (String image : dockerCmdExecFactory.getImageNames()) {
            LOG.info("Cleaning up temporary image with {}", image);
            try {
                dockerClient.removeImageCmd(image).withForce(true).exec();
            } catch (DockerException ignore) {
                // ignore.printStackTrace();
            }
        }

        for (String volume : dockerCmdExecFactory.getVolumeNames()) {
            LOG.info("Cleaning up temporary volume with {}", volume);
            try {
                dockerClient.removeVolumeCmd(volume).exec();
            } catch (DockerException ignore) {
                // ignore.printStackTrace();
            }
        }

        for (String networkId : dockerCmdExecFactory.getNetworkIds()) {
            LOG.info("Cleaning up temporary network with {}", networkId);
            try {
                dockerClient.removeNetworkCmd(networkId).exec();
            } catch (DockerException ignore) {
                // ignore.printStackTrace();
            }
        }

        LOG.info("################################## END OF {} ##################################\n", result.getName());
    }

    protected String asString(InputStream response) {
        return consumeAsString(response);
    }

    public static String consumeAsString(InputStream response) {

        StringWriter logwriter = new StringWriter();

        try {
            LineIterator itr = IOUtils.lineIterator(response, "UTF-8");

            while (itr.hasNext()) {
                String line = itr.next();
                logwriter.write(line + (itr.hasNext() ? "\n" : ""));
                LOG.info("line: " + line);
            }
            response.close();

            return logwriter.toString();
        } catch (IOException e) {
            throw new RuntimeException(e);
        } finally {
            IOUtils.closeQuietly(response);
        }
    }

    // UTIL

    /**
     * Checks to see if a specific port is available.
     *
     * @param port
     *            the port to check for availability
     */
    public static Boolean available(int port) {
        if (port < 1100 || port > 60000) {
            throw new IllegalArgumentException("Invalid start port: " + port);
        }

        ServerSocket ss = null;
        DatagramSocket ds = null;
        try {
            ss = new ServerSocket(port);
            ss.setReuseAddress(true);
            ds = new DatagramSocket(port);
            ds.setReuseAddress(true);
            return true;
        } catch (IOException ignored) {
        } finally {
            if (ds != null) {
                ds.close();
            }

            if (ss != null) {
                try {
                    ss.close();
                } catch (IOException e) {
                    /* should not be thrown */
                }
            }
        }

        return false;
    }

    protected MountedVolumes mountedVolumes(Matcher<? super List<Volume>> subMatcher) {
        return new MountedVolumes(subMatcher, "Mounted volumes", "mountedVolumes");
    }

    private static class MountedVolumes extends FeatureMatcher<InspectContainerResponse, List<Volume>> {
        public MountedVolumes(Matcher<? super List<Volume>> subMatcher, String featureDescription, String featureName) {
            super(subMatcher, featureDescription, featureName);
        }

        @Override
        public List<Volume> featureValueOf(InspectContainerResponse item) {
            List<Volume> volumes = new ArrayList<Volume>();
            for (Mount mount : item.getMounts()) {
                volumes.add(mount.getDestination());
            }
            return volumes;
        }
    }

    protected String containerLog(String containerId) throws Exception {
        return dockerClient.logContainerCmd(containerId).withStdOut(true).exec(new LogContainerTestCallback())
                .awaitCompletion().toString();
    }

    public static class LogContainerTestCallback extends LogContainerResultCallback {
        protected final StringBuffer log = new StringBuffer();

        List<Frame> collectedFrames = new ArrayList<Frame>();

        boolean collectFrames = false;

        public LogContainerTestCallback() {
            this(false);
        }

        public LogContainerTestCallback(boolean collectFrames) {
            this.collectFrames = collectFrames;
        }

        @Override
        public void onNext(Frame frame) {
            if(collectFrames) collectedFrames.add(frame);
            log.append(new String(frame.getPayload()));
        }

        @Override
        public String toString() {
            return log.toString();
        }


        public List<Frame> getCollectedFrames() {
            return collectedFrames;
        }
    }

    protected String buildImage(File baseDir) throws Exception {

        return dockerClient.buildImageCmd(baseDir).withNoCache(true).exec(new BuildImageResultCallback())
                .awaitImageId();
    }

    protected Network findNetwork(List<Network> networks, String name) {

        for (Network network : networks) {
            if (StringUtils.equals(network.getName(), name)) {
                return network;
            }
        }

        fail("No network found.");
        return null;
    }

}