Added withLogConsumer() method to DockerComposeContainer (#629)

* Added withLogConsumer() method to DockerComposeContainer

Fixes #605

* Added try-finally block to DockerComposeLogConsumerTest

Author: barrycommins

CHANGELOG.md

@@ -19,6 +19,7 @@ All notable changes to this project will be documented in this file.
 - Added `WaitStrategyTarget` interface which is the target of the new `WaitStrategy` ([\#600](https://github.com/testcontainers/testcontainers-java/pull/600))
 - *Breaking:* Removed hard-coded `wnameless` Oracle database image name. Users should instead place a file on the classpath named `testcontainers.properties` containing `oracle.container.image=IMAGE`, where IMAGE is a suitable image name and tag/SHA hash. For information, the approach recommended by Oracle for creating an Oracle XE docker image is described [here](https://blogs.oracle.com/oraclewebcentersuite/implement-oracle-database-xe-as-docker-containers). 
 - Added `DockerHealthcheckWaitStrategy` that is based on Docker's built-in [healthcheck](https://docs.docker.com/engine/reference/builder/#healthcheck) ([\#618](https://github.com/testcontainers/testcontainers-java/pull/618)).
+- Added `withLogConsumer(String serviceName, Consumer<OutputFrame> consumer)` method to `DockerComposeContainer` ([\#605](https://github.com/testcontainers/testcontainers-java/issues/605))
 
 ## [1.6.0] - 2018-01-28
 

core/src/main/java/org/testcontainers/containers/DockerComposeContainer.java

@@ -15,6 +15,7 @@
 import org.slf4j.LoggerFactory;
 import org.slf4j.profiler.Profiler;
 import org.testcontainers.DockerClientFactory;
+import org.testcontainers.containers.output.OutputFrame;
 import org.testcontainers.containers.output.Slf4jLogConsumer;
 import org.testcontainers.containers.startupcheck.IndefiniteWaitOneShotStartupCheckStrategy;
 import org.testcontainers.containers.wait.strategy.Wait;
@@ -31,7 +32,9 @@
 import java.nio.file.Paths;
 import java.time.Duration;
 import java.util.AbstractMap.SimpleEntry;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -40,6 +43,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
@@ -72,6 +76,7 @@ public class DockerComposeContainer<SELF extends DockerComposeContainer<SELF>> e
     private final Map<String, ComposeServiceWaitStrategyTarget> serviceInstanceMap = new ConcurrentHashMap<>();
     private final Map<String, WaitAllStrategy> waitStrategyMap = new ConcurrentHashMap<>();
     private final SocatContainer ambassadorContainer = new SocatContainer();
+    private final Map<String, List<Consumer<OutputFrame>>> logConsumers = new ConcurrentHashMap<>();
 
     private static final Object MUTEX = new Object();
 
@@ -148,11 +153,13 @@ private void createServiceInstance(Container container) {
         final ComposeServiceWaitStrategyTarget containerInstance = new ComposeServiceWaitStrategyTarget(container,
             ambassadorContainer, ambassadorPortMappings.getOrDefault(serviceName, new HashMap<>()));
 
+        String containerId = containerInstance.getContainerId();
         if (tailChildContainers) {
-            LogUtils.followOutput(DockerClientFactory.instance().client(), containerInstance.getContainerId(),
-                new Slf4jLogConsumer(logger()).withPrefix(container.getNames()[0]));
+            followLogs(containerId, new Slf4jLogConsumer(logger()).withPrefix(container.getNames()[0]));
         }
-       serviceInstanceMap.putIfAbsent(serviceName, containerInstance);
+        //follow logs using registered consumers for this service
+        logConsumers.getOrDefault(serviceName, Collections.emptyList()).forEach(consumer -> followLogs(containerId, consumer));
+        serviceInstanceMap.putIfAbsent(serviceName, containerInstance);
     }
 
     private void waitUntilServiceStarted(String serviceName, ComposeServiceWaitStrategyTarget serviceInstance) {
@@ -401,6 +408,27 @@ public SELF withTailChildContainers(boolean tailChildContainers) {
         return self();
     }
 
+    /**
+     * Attach an output consumer at container startup, enabling stdout and stderr to be followed, waited on, etc.
+     * <p>
+     * More than one consumer may be registered.
+     *
+     * @param serviceName the name of the service as set in the docker-compose.yml file
+     * @param consumer consumer that output frames should be sent to
+     * @return this instance, for chaining
+     */
+    public SELF withLogConsumer(String serviceName, Consumer<OutputFrame> consumer) {
+        String serviceInstanceName = getServiceInstanceName(serviceName);
+        final List<Consumer<OutputFrame>> consumers = this.logConsumers.getOrDefault(serviceInstanceName, new ArrayList<>());
+        consumers.add(consumer);
+        this.logConsumers.putIfAbsent(serviceInstanceName, consumers);
+        return self();
+    }
+
+    private void followLogs(String containerId, Consumer<OutputFrame> consumer) {
+        LogUtils.followOutput(DockerClientFactory.instance().client(), containerId, consumer);
+    }
+
     private SELF self() {
         return (SELF) this;
     }

core/src/test/java/org/testcontainers/junit/DockerComposeLogConsumerTest.java

@@ -0,0 +1,30 @@
+package org.testcontainers.junit;
+
+import org.junit.Test;
+import org.junit.runner.Description;
+import org.testcontainers.containers.DockerComposeContainer;
+import org.testcontainers.containers.output.WaitingConsumer;
+
+import java.io.File;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import static org.testcontainers.containers.output.OutputFrame.OutputType.STDOUT;
+
+public class DockerComposeLogConsumerTest {
+
+    @Test
+    public void testLogConsumer() throws TimeoutException {
+        WaitingConsumer logConsumer = new WaitingConsumer();
+        DockerComposeContainer environment = new DockerComposeContainer(new File("src/test/resources/v2-compose-test.yml"))
+            .withExposedService("redis_1", 6379)
+            .withLogConsumer("redis_1", logConsumer);
+
+        try {
+            environment.starting(Description.EMPTY);
+            logConsumer.waitUntil(frame -> frame.getType() == STDOUT && frame.getUtf8String().contains("Ready to accept connections"), 5, TimeUnit.SECONDS);
+        } finally {
+            environment.finished(Description.EMPTY);
+        }
+    }
+}