-
-
Notifications
You must be signed in to change notification settings - Fork 1.7k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
3 changed files
with
175 additions
and
12 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
87 changes: 87 additions & 0 deletions
87
modules/kafka/src/main/java/org/testcontainers/containers/KafkaContainerCluster.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,87 @@ | ||
package org.testcontainers.containers; | ||
|
||
import java.util.Collection; | ||
import java.util.stream.Collectors; | ||
import java.util.stream.IntStream; | ||
import java.util.stream.Stream; | ||
|
||
import static org.testcontainers.containers.KafkaContainer.CONFLUENT_PLATFORM_VERSION; | ||
|
||
/** | ||
* Provides an easy way to launch a Kafka cluster with multiple brokers. | ||
*/ | ||
public class KafkaContainerCluster implements AutoCloseable { | ||
|
||
private final int startPort; | ||
private final Network network; | ||
private final GenericContainer zookeeper; | ||
private final Collection<KafkaContainer> brokers; | ||
|
||
public KafkaContainerCluster(int brokersNum, int internalTopicsRf, int startPort) { | ||
this(CONFLUENT_PLATFORM_VERSION, brokersNum, internalTopicsRf, startPort); | ||
} | ||
|
||
public KafkaContainerCluster(String confluentPlatformVersion, int brokersNum, int internalTopicsRf, int startPort) { | ||
if (brokersNum < 0) { | ||
throw new IllegalArgumentException("brokersNum '" + brokersNum + "' must be greater than 0"); | ||
} | ||
if (internalTopicsRf < 0 || internalTopicsRf > brokersNum) { | ||
throw new IllegalArgumentException("internalTopicsRf '" + internalTopicsRf + "' must be less than brokersNum and greater than 0"); | ||
} | ||
|
||
this.startPort = startPort; | ||
|
||
this.network = Network.newNetwork(); | ||
|
||
this.zookeeper = new GenericContainer("confluentinc/cp-zookeeper:" + confluentPlatformVersion) | ||
.withNetwork(network) | ||
.withNetworkAliases("zookeeper") | ||
.withEnv("ZOOKEEPER_CLIENT_PORT", String.valueOf(KafkaContainer.ZOOKEEPER_PORT)); | ||
|
||
this.brokers = IntStream | ||
.range(0, brokersNum) | ||
.mapToObj(brokerNum -> | ||
new KafkaContainer(confluentPlatformVersion, this.startPort + brokerNum, String.valueOf(brokerNum), internalTopicsRf) | ||
.withNetwork(this.network) | ||
.withExternalZookeeper("zookeeper:" + KafkaContainer.ZOOKEEPER_PORT) | ||
) | ||
.collect(Collectors.toList()); | ||
} | ||
|
||
public Network getNetwork() { | ||
return this.network; | ||
} | ||
|
||
public GenericContainer getZooKeeper() { | ||
return this.zookeeper; | ||
} | ||
|
||
public Collection<KafkaContainer> getBrokers() { | ||
return this.brokers; | ||
} | ||
|
||
public String getBootstrapServers() { | ||
return brokers.stream() | ||
.map(KafkaContainer::getBootstrapServers) | ||
.collect(Collectors.joining(",")); | ||
} | ||
|
||
private Stream<GenericContainer> allContainers() { | ||
Stream<GenericContainer> genericBrokers = this.brokers.stream().map(b -> (GenericContainer) b); | ||
Stream<GenericContainer> zookeeper = Stream.of(this.zookeeper); | ||
return Stream.concat(genericBrokers, zookeeper); | ||
} | ||
|
||
public void startAll() { | ||
allContainers().parallel().forEach(GenericContainer::start); | ||
} | ||
|
||
public void stopAll() { | ||
allContainers().parallel().forEach(GenericContainer::stop); | ||
} | ||
|
||
@Override | ||
public void close() throws Exception { | ||
this.stopAll(); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters