Repository: kafka
Updated Branches:
  refs/heads/trunk b938c03b0 -> 7837d3e54


KAFKA-3896: Fix KStreamRepartitionJoinTest

The root cause of this issue is that in InternalTopicManager we are creating 
topics one-at-a-time, and for this test, there are 31 topics to be created, as 
a result it is possible that the consumer could time out during the assignment 
in rebalance, and the next leader has to do the same again because of 
"makeReady" calls are one-at-a-time.

This patch batches the topics into a single create request and also use the 
StreamsKafkaClient directly to fetch metadata for validating the created 
topics. Also optimized a bunch of inefficient code in InternalTopicManager and 
StreamsKafkaClient.

Minor cleanup: make the exception message more informative in integration tests.

Author: Guozhang Wang <[email protected]>

Reviewers: Damian Guy, Matthias J. Sax, Jason Gustafson

Closes #2405 from guozhangwang/K3896-fix-kstream-repartition-join-test


Project: http://git-wip-us.apache.org/repos/asf/kafka/repo
Commit: http://git-wip-us.apache.org/repos/asf/kafka/commit/7837d3e5
Tree: http://git-wip-us.apache.org/repos/asf/kafka/tree/7837d3e5
Diff: http://git-wip-us.apache.org/repos/asf/kafka/diff/7837d3e5

Branch: refs/heads/trunk
Commit: 7837d3e54850fc6ecd16f9bc4dd28a02779572c2
Parents: b938c03
Author: Guozhang Wang <[email protected]>
Authored: Tue Jan 24 12:00:35 2017 -0800
Committer: Guozhang Wang <[email protected]>
Committed: Tue Jan 24 12:00:35 2017 -0800

----------------------------------------------------------------------
 .../internals/InternalTopicManager.java         | 86 +++++++++-----------
 .../internals/StreamPartitionAssignor.java      | 42 +++++++---
 .../processor/internals/StreamsKafkaClient.java | 40 +--------
 .../integration/KStreamRepartitionJoinTest.java |  8 +-
 .../integration/utils/IntegrationTestUtils.java |  4 +-
 .../internals/InternalTopicManagerTest.java     | 66 ++++++---------
 .../kafka/test/MockInternalTopicManager.java    | 28 +++++--
 7 files changed, 123 insertions(+), 151 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/kafka/blob/7837d3e5/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopicManager.java
----------------------------------------------------------------------
diff --git 
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopicManager.java
 
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopicManager.java
index 133375f..eeb38d1 100644
--- 
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopicManager.java
+++ 
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopicManager.java
@@ -26,6 +26,7 @@ import java.io.IOException;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
 public class InternalTopicManager {
@@ -48,23 +49,17 @@ public class InternalTopicManager {
     }
 
     /**
-     * Prepares a given internal topic.
-     * If the topic does not exist creates a new topic.
-     * If the topic with the correct number of partitions exists ignores it.
-     * If the topic exists already but has different number of partitions we 
fail and throw exception requesting user to reset the app before restarting 
again.
+     * Prepares a set of given internal topics.
      *
-     * @param topic
-     * @param numPartitions
+     * If a topic does not exist creates a new topic.
+     * If a topic with the correct number of partitions exists ignores it.
+     * If a topic exists already but has different number of partitions we 
fail and throw exception requesting user to reset the app before restarting 
again.
      */
-    public void makeReady(final InternalTopicConfig topic, int numPartitions) {
-
-        Map<InternalTopicConfig, Integer> topics = new HashMap<>();
-        topics.put(topic, numPartitions);
+    public void makeReady(final Map<InternalTopicConfig, Integer> topics) {
         for (int i = 0; i < MAX_TOPIC_READY_TRY; i++) {
             try {
-                Collection<MetadataResponse.TopicMetadata> topicsMetadata = 
streamsKafkaClient.fetchTopicsMetadata();
-                validateTopicPartitons(topics, topicsMetadata);
-                Map<InternalTopicConfig, Integer> topicsToBeCreated = 
filterExistingTopics(topics, topicsMetadata);
+                final Map<String, Integer> existingTopicPartitions = 
fetchExistingPartitionCountByTopic();
+                final Map<InternalTopicConfig, Integer> topicsToBeCreated = 
validateTopicPartitions(topics, existingTopicPartitions);
                 streamsKafkaClient.createTopics(topicsToBeCreated, 
replicationFactor, windowChangeLogAdditionalRetention);
                 return;
             } catch (StreamsException ex) {
@@ -74,6 +69,16 @@ public class InternalTopicManager {
         throw new StreamsException("Could not create internal topics.");
     }
 
+    /**
+     * Get the number of partitions for the given topics
+     */
+    public Map<String, Integer> getNumPartitions(final Set<String> topics) {
+        final Map<String, Integer> existingTopicPartitions = 
fetchExistingPartitionCountByTopic();
+        existingTopicPartitions.keySet().retainAll(topics);
+
+        return existingTopicPartitions;
+    }
+
     public void close() {
         try {
             streamsKafkaClient.close();
@@ -83,49 +88,36 @@ public class InternalTopicManager {
     }
 
     /**
-     * Return the non existing topics.
-     *
-     * @param topicsPartitionsMap
-     * @param topicsMetadata
-     * @return
+     * Check the existing topics to have correct number of partitions; and 
return the non existing topics to be created
      */
-    private Map<InternalTopicConfig, Integer> filterExistingTopics(final 
Map<InternalTopicConfig, Integer> topicsPartitionsMap, 
Collection<MetadataResponse.TopicMetadata> topicsMetadata) {
-        Map<String, Integer> existingTopicNamesPartitions = 
getExistingTopicNamesPartitions(topicsMetadata);
-        Map<InternalTopicConfig, Integer> nonExistingTopics = new HashMap<>();
-        // Add the topics that don't exist to the nonExistingTopics.
+    private Map<InternalTopicConfig, Integer> validateTopicPartitions(final 
Map<InternalTopicConfig, Integer> topicsPartitionsMap,
+                                                                      final 
Map<String, Integer> existingTopicNamesPartitions) {
+        final Map<InternalTopicConfig, Integer> topicsToBeCreated = new 
HashMap<>();
         for (InternalTopicConfig topic: topicsPartitionsMap.keySet()) {
-            if (existingTopicNamesPartitions.get(topic.name()) == null) {
-                nonExistingTopics.put(topic, topicsPartitionsMap.get(topic));
+            if (existingTopicNamesPartitions.containsKey(topic.name())) {
+                if 
(!existingTopicNamesPartitions.get(topic.name()).equals(topicsPartitionsMap.get(topic)))
 {
+                    throw new StreamsException("Existing internal topic " + 
topic.name() + " has invalid partitions." +
+                            " Expected: " + topicsPartitionsMap.get(topic) + " 
Actual: " + existingTopicNamesPartitions.get(topic.name()) +
+                            ". Use 'kafka.tools.StreamsResetter' tool to clean 
up invalid topics before processing.");
+                }
+            } else {
+                topicsToBeCreated.put(topic, topicsPartitionsMap.get(topic));
             }
         }
-        return nonExistingTopics;
+
+        return topicsToBeCreated;
     }
 
+    private Map<String, Integer> fetchExistingPartitionCountByTopic() {
+        // The names of existing topics and corresponding partition counts
+        final Map<String, Integer> existingPartitionCountByTopic = new 
HashMap<>();
 
-    /**
-     * Make sure the existing topics have correct number of partitions.
-     *
-     * @param topicsPartitionsMap
-     * @param topicsMetadata
-     */
-    private void validateTopicPartitons(final Map<InternalTopicConfig, 
Integer> topicsPartitionsMap, Collection<MetadataResponse.TopicMetadata> 
topicsMetadata) {
-        Map<String, Integer> existingTopicNamesPartitions = 
getExistingTopicNamesPartitions(topicsMetadata);
-        for (InternalTopicConfig topic: topicsPartitionsMap.keySet()) {
-            if (existingTopicNamesPartitions.get(topic.name()) != null) {
-                if (existingTopicNamesPartitions.get(topic.name()) != 
topicsPartitionsMap.get(topic)) {
-                    throw new StreamsException("Internal topic with invalid 
partitons. Use 'kafka.tools.StreamsResetter' tool to clean up invalid topics 
before processing.");
-                }
-            }
-        }
-    }
+        Collection<MetadataResponse.TopicMetadata> topicsMetadata = 
streamsKafkaClient.fetchTopicsMetadata();
 
-    private Map<String, Integer> 
getExistingTopicNamesPartitions(Collection<MetadataResponse.TopicMetadata> 
topicsMetadata) {
-        // The names of existing topics
-        Map<String, Integer> existingTopicNamesPartitions = new HashMap<>();
         for (MetadataResponse.TopicMetadata topicMetadata: topicsMetadata) {
-            existingTopicNamesPartitions.put(topicMetadata.topic(), 
topicMetadata.partitionMetadata().size());
+            existingPartitionCountByTopic.put(topicMetadata.topic(), 
topicMetadata.partitionMetadata().size());
         }
-        return existingTopicNamesPartitions;
-    }
 
+        return existingPartitionCountByTopic;
+    }
 }

http://git-wip-us.apache.org/repos/asf/kafka/blob/7837d3e5/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamPartitionAssignor.java
----------------------------------------------------------------------
diff --git 
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamPartitionAssignor.java
 
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamPartitionAssignor.java
index 382c4e9..c790a1e 100644
--- 
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamPartitionAssignor.java
+++ 
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamPartitionAssignor.java
@@ -263,7 +263,6 @@ public class StreamPartitionAssignor implements 
PartitionAssignor, Configurable
      */
     @Override
     public Map<String, Assignment> assign(Cluster metadata, Map<String, 
Subscription> subscriptions) {
-
         // construct the client metadata from the decoded subscription info
         Map<UUID, ClientMetadata> clientsMetadata = new HashMap<>();
 
@@ -593,9 +592,13 @@ public class StreamPartitionAssignor implements 
PartitionAssignor, Configurable
     private void prepareTopic(final Map<String, InternalTopicMetadata> 
topicPartitions) {
         log.debug("stream-thread [{}] Starting to validate internal topics in 
partition assignor.", streamThread.getName());
 
-        for (final Map.Entry<String, InternalTopicMetadata> entry : 
topicPartitions.entrySet()) {
-            final InternalTopicConfig topic = entry.getValue().config;
-            final Integer numPartitions = entry.getValue().numPartitions;
+        // first construct the topics to make ready
+        Map<InternalTopicConfig, Integer> topicsToMakeReady = new HashMap<>();
+        Set<String> topicNamesToMakeReady = new HashSet<>();
+
+        for (InternalTopicMetadata metadata : topicPartitions.values()) {
+            InternalTopicConfig topic = metadata.config;
+            Integer numPartitions = metadata.numPartitions;
 
             if (numPartitions == NOT_AVAILABLE) {
                 continue;
@@ -604,18 +607,37 @@ public class StreamPartitionAssignor implements 
PartitionAssignor, Configurable
                 throw new 
TopologyBuilderException(String.format("stream-thread [%s] Topic [%s] number of 
partitions not defined", streamThread.getName(), topic.name()));
             }
 
-            internalTopicManager.makeReady(topic, numPartitions);
+            topicsToMakeReady.put(topic, numPartitions);
+            topicNamesToMakeReady.add(topic.name());
+        }
 
-            // wait until the topic metadata has been propagated to all brokers
-            List<PartitionInfo> partitions;
-            do {
-                partitions = 
streamThread.restoreConsumer.partitionsFor(topic.name());
-            } while (partitions == null || partitions.size() != numPartitions);
+        if (!topicsToMakeReady.isEmpty()) {
+            internalTopicManager.makeReady(topicsToMakeReady);
+
+            // wait until each one of the topic metadata has been propagated 
to at least one broker
+            while (!allTopicsCreated(topicNamesToMakeReady, 
topicsToMakeReady)) {
+                try {
+                    Thread.sleep(50L);
+                } catch (InterruptedException e) {
+                    // ignore
+                }
+            }
         }
 
         log.info("stream-thread [{}] Completed validating internal topics in 
partition assignor", streamThread.getName());
     }
 
+    private boolean allTopicsCreated(final Set<String> topicNamesToMakeReady, 
final Map<InternalTopicConfig, Integer> topicsToMakeReady) {
+        final Map<String, Integer> partitions = 
internalTopicManager.getNumPartitions(topicNamesToMakeReady);
+        for (Map.Entry<InternalTopicConfig, Integer> entry : 
topicsToMakeReady.entrySet()) {
+            final Integer numPartitions = 
partitions.get(entry.getKey().name());
+            if (numPartitions == null || 
!numPartitions.equals(entry.getValue())) {
+                return false;
+            }
+        }
+        return true;
+    }
+
     private void ensureCopartitioning(Collection<Set<String>> 
copartitionGroups,
                                       Map<String, InternalTopicMetadata> 
allRepartitionTopicsNumPartitions,
                                       Cluster metadata) {

http://git-wip-us.apache.org/repos/asf/kafka/blob/7837d3e5/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsKafkaClient.java
----------------------------------------------------------------------
diff --git 
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsKafkaClient.java
 
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsKafkaClient.java
index 8ea570f..df590fb 100644
--- 
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsKafkaClient.java
+++ 
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsKafkaClient.java
@@ -48,7 +48,6 @@ import java.util.LinkedHashMap;
 import java.util.Properties;
 import java.util.HashMap;
 import java.util.Collection;
-import java.util.Arrays;
 import java.util.concurrent.TimeUnit;
 
 public class StreamsKafkaClient {
@@ -103,10 +102,6 @@ public class StreamsKafkaClient {
 
     /**
      * Creates a set of new topics using batch request.
-     *
-     * @param topicsMap
-     * @param replicationFactor
-     * @param windowChangeLogAdditionalRetention
      */
     public void createTopics(final Map<InternalTopicConfig, Integer> 
topicsMap, final int replicationFactor, final long 
windowChangeLogAdditionalRetention) {
 
@@ -135,12 +130,8 @@ public class StreamsKafkaClient {
 
         for (InternalTopicConfig internalTopicConfig : topicsMap.keySet()) {
             CreateTopicsResponse.Error error = 
createTopicsResponse.errors().get(internalTopicConfig.name());
-            if (!error.is(Errors.NONE)) {
-                if (error.is(Errors.TOPIC_ALREADY_EXISTS)) {
-                    continue;
-                } else {
-                    throw new StreamsException("Could not create topic: " + 
internalTopicConfig.name() + " due to " + error.messageWithFallback());
-                }
+            if (!error.is(Errors.NONE) && 
!error.is(Errors.TOPIC_ALREADY_EXISTS)) {
+                throw new StreamsException("Could not create topic: " + 
internalTopicConfig.name() + " due to " + error.messageWithFallback());
             }
         }
     }
@@ -193,35 +184,8 @@ public class StreamsKafkaClient {
 
     }
 
-
-     /**
-     * Fetch the metadata for a topic.
-     * @param topic
-     * @return
-     */
-    public MetadataResponse.TopicMetadata fetchTopicMetadata(final String 
topic) {
-        final ClientRequest clientRequest = 
kafkaClient.newClientRequest(getBrokerId(), new 
MetadataRequest.Builder(Arrays.asList(topic)), Time.SYSTEM.milliseconds(), 
true, null);
-        final ClientResponse clientResponse = sendRequest(clientRequest);
-        if (!clientResponse.hasResponse()) {
-            throw new StreamsException("Empty response for client request.");
-        }
-        if (!(clientResponse.responseBody() instanceof MetadataResponse)) {
-            throw new StreamsException("Inconsistent response type for 
internal topic metadata request. Expected MetadataResponse but received " + 
clientResponse.responseBody().getClass().getName());
-        }
-        final MetadataResponse metadataResponse = (MetadataResponse) 
clientResponse.responseBody();
-        for (MetadataResponse.TopicMetadata topicMetadata: 
metadataResponse.topicMetadata()) {
-            if (topicMetadata.topic().equalsIgnoreCase(topic)) {
-                return topicMetadata;
-            }
-        }
-        return null;
-    }
-
-
     /**
      * Fetch the metadata for all topics
-     *
-     * @return
      */
     public Collection<MetadataResponse.TopicMetadata> fetchTopicsMetadata() {
 

http://git-wip-us.apache.org/repos/asf/kafka/blob/7837d3e5/streams/src/test/java/org/apache/kafka/streams/integration/KStreamRepartitionJoinTest.java
----------------------------------------------------------------------
diff --git 
a/streams/src/test/java/org/apache/kafka/streams/integration/KStreamRepartitionJoinTest.java
 
b/streams/src/test/java/org/apache/kafka/streams/integration/KStreamRepartitionJoinTest.java
index f0fb0a2..43e5d87 100644
--- 
a/streams/src/test/java/org/apache/kafka/streams/integration/KStreamRepartitionJoinTest.java
+++ 
b/streams/src/test/java/org/apache/kafka/streams/integration/KStreamRepartitionJoinTest.java
@@ -160,7 +160,6 @@ public class KStreamRepartitionJoinTest {
         return new ExpectedOutputOnTopic(expectedStreamOneTwoJoin, 
"map-both-streams-and-join-" + testNo);
     }
 
-
     private ExpectedOutputOnTopic mapMapJoin() throws Exception {
         final KStream<Integer, Integer> mapMapStream = streamOne.map(
             new KeyValueMapper<Long, Integer, KeyValue<Long, Integer>>() {
@@ -178,8 +177,7 @@ public class KStreamRepartitionJoinTest {
         return new ExpectedOutputOnTopic(expectedStreamOneTwoJoin, 
outputTopic);
     }
 
-
-    public ExpectedOutputOnTopic selectKeyAndJoin() throws ExecutionException, 
InterruptedException {
+    private ExpectedOutputOnTopic selectKeyAndJoin() throws 
ExecutionException, InterruptedException {
 
         final KStream<Integer, Integer> keySelected =
             streamOne.selectKey(MockKeyValueMapper.<Long, 
Integer>SelectValueMapper());
@@ -189,7 +187,6 @@ public class KStreamRepartitionJoinTest {
         return new ExpectedOutputOnTopic(expectedStreamOneTwoJoin, 
outputTopic);
     }
 
-
     private ExpectedOutputOnTopic flatMapJoin() throws Exception {
         final KStream<Integer, Integer> flatMapped = streamOne.flatMap(
             new KeyValueMapper<Long, Integer, Iterable<KeyValue<Integer, 
Integer>>>() {
@@ -221,7 +218,7 @@ public class KStreamRepartitionJoinTest {
         return new ExpectedOutputOnTopic(Arrays.asList("A:1", "B:2", "C:3", 
"D:4", "E:5"), output);
     }
 
-    public ExpectedOutputOnTopic mapBothStreamsAndLeftJoin() throws Exception {
+    private ExpectedOutputOnTopic mapBothStreamsAndLeftJoin() throws Exception 
{
         final KStream<Integer, Integer> map1 = streamOne.map(keyMapper);
 
         final KStream<Integer, String> map2 = 
streamTwo.map(MockKeyValueMapper.<Integer, String>NoOpKeyValueMapper());
@@ -312,7 +309,6 @@ public class KStreamRepartitionJoinTest {
 
     }
 
-
     private void produceStreamTwoInputTo(final String streamTwoInput)
         throws ExecutionException, InterruptedException {
         IntegrationTestUtils.produceKeyValuesSynchronously(

http://git-wip-us.apache.org/repos/asf/kafka/blob/7837d3e5/streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java
----------------------------------------------------------------------
diff --git 
a/streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java
 
b/streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java
index 3e7c41b..08e22cc 100644
--- 
a/streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java
+++ 
b/streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java
@@ -214,7 +214,7 @@ public class IntegrationTestUtils {
             }
         };
 
-        final String conditionDetails = "Expecting " + expectedNumRecords + " 
records while only received " + accumData.size() + ": " + accumData;
+        final String conditionDetails = "Expecting " + expectedNumRecords + " 
records from topic " + topic + " while only received " + accumData.size() + ": 
" + accumData;
 
         TestUtils.waitForCondition(valuesRead, waitTime, conditionDetails);
 
@@ -254,7 +254,7 @@ public class IntegrationTestUtils {
             }
         };
 
-        final String conditionDetails = "Expecting " + expectedNumRecords + " 
records while only received " + accumData.size() + ": " + accumData;
+        final String conditionDetails = "Expecting " + expectedNumRecords + " 
records from topic " + topic + " while only received " + accumData.size() + ": 
" + accumData;
 
         TestUtils.waitForCondition(valuesRead, waitTime, conditionDetails);
 

http://git-wip-us.apache.org/repos/asf/kafka/blob/7837d3e5/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopicManagerTest.java
----------------------------------------------------------------------
diff --git 
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopicManagerTest.java
 
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopicManagerTest.java
index f828099..548cb45 100644
--- 
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopicManagerTest.java
+++ 
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/InternalTopicManagerTest.java
@@ -23,47 +23,55 @@ import org.apache.kafka.common.requests.MetadataResponse;
 import org.apache.kafka.streams.StreamsConfig;
 import org.apache.kafka.streams.errors.StreamsException;
 import org.apache.kafka.test.MockTimestampExtractor;
+import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
+import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Map;
 import java.util.Properties;
-import java.util.Arrays;
-import java.util.ArrayList;
 
 import static 
org.apache.kafka.streams.processor.internals.InternalTopicManager.WINDOW_CHANGE_LOG_ADDITIONAL_RETENTION_DEFAULT;
 
 public class InternalTopicManagerTest {
 
-    private String userEndPoint = "localhost:2171";
-    StreamsConfig config;
-    MockStreamKafkaClient streamsKafkaClient;
+    private final String topic = "test_topic";
+    private final String userEndPoint = "localhost:2171";
+    private MockStreamKafkaClient streamsKafkaClient;
 
     @Before
     public void init() {
-        config = new StreamsConfig(configProps());
+        final StreamsConfig config = new StreamsConfig(configProps());
         streamsKafkaClient = new MockStreamKafkaClient(config);
     }
 
+    @After
+    public void shutdown() throws IOException {
+        streamsKafkaClient.close();
+    }
+
     @Test
-    public void shouldCreateRequiredTopics() throws Exception {
+    public void shouldReturnCorrectPartitionCounts() throws Exception {
+        InternalTopicManager internalTopicManager = new 
InternalTopicManager(streamsKafkaClient, 1, 
WINDOW_CHANGE_LOG_ADDITIONAL_RETENTION_DEFAULT);
+        Assert.assertEquals(Collections.singletonMap(topic, 1), 
internalTopicManager.getNumPartitions(Collections.singleton(topic)));
+    }
 
-        streamsKafkaClient.setReturnCorrectTopic(true);
+    @Test
+    public void shouldCreateRequiredTopics() throws Exception {
         InternalTopicManager internalTopicManager = new 
InternalTopicManager(streamsKafkaClient, 1, 
WINDOW_CHANGE_LOG_ADDITIONAL_RETENTION_DEFAULT);
-        internalTopicManager.makeReady(new InternalTopicConfig("test_topic", 
Collections.singleton(InternalTopicConfig.CleanupPolicy.compact), null), 1);
+        internalTopicManager.makeReady(Collections.singletonMap(new 
InternalTopicConfig(topic, 
Collections.singleton(InternalTopicConfig.CleanupPolicy.compact), null), 1));
     }
 
     @Test
     public void shouldNotCreateTopicIfExistsWithDifferentPartitions() throws 
Exception {
-
-        streamsKafkaClient.setReturnCorrectTopic(true);
         InternalTopicManager internalTopicManager = new 
InternalTopicManager(streamsKafkaClient, 1, 
WINDOW_CHANGE_LOG_ADDITIONAL_RETENTION_DEFAULT);
         boolean exceptionWasThrown = false;
         try {
-            internalTopicManager.makeReady(new 
InternalTopicConfig("test_topic", 
Collections.singleton(InternalTopicConfig.CleanupPolicy.compact), null), 2);
+            internalTopicManager.makeReady(Collections.singletonMap(new 
InternalTopicConfig(topic, 
Collections.singleton(InternalTopicConfig.CleanupPolicy.compact), null), 2));
         } catch (StreamsException e) {
             exceptionWasThrown = true;
         }
@@ -82,43 +90,21 @@ public class InternalTopicManagerTest {
     }
 
     private class MockStreamKafkaClient extends StreamsKafkaClient {
-        public MockStreamKafkaClient(final StreamsConfig streamsConfig) {
-            super(streamsConfig);
-        }
 
-        public boolean isReturnCorrectTopic() {
-            return returnCorrectTopic;
-        }
-
-        public void setReturnCorrectTopic(boolean returnCorrectTopic) {
-            this.returnCorrectTopic = returnCorrectTopic;
+        MockStreamKafkaClient(final StreamsConfig streamsConfig) {
+            super(streamsConfig);
         }
 
-        boolean returnCorrectTopic = false;
-
-
         @Override
         public void createTopics(final Map<InternalTopicConfig, Integer> 
topicsMap, final int replicationFactor, final long 
windowChangeLogAdditionalRetention) {
-
-        }
-
-        @Override
-        public MetadataResponse.TopicMetadata fetchTopicMetadata(final String 
topic) {
-
-            if (returnCorrectTopic) {
-                MetadataResponse.PartitionMetadata partitionMetadata = new 
MetadataResponse.PartitionMetadata(Errors.NONE, 1, null, new ArrayList<Node>(), 
new ArrayList<Node>());
-                MetadataResponse.TopicMetadata topicMetadata = new 
MetadataResponse.TopicMetadata(Errors.NONE, topic, true, 
Arrays.asList(partitionMetadata));
-                return topicMetadata;
-            }
-            return null;
+            // do nothing
         }
 
         @Override
         public Collection<MetadataResponse.TopicMetadata> 
fetchTopicsMetadata() {
-            if (returnCorrectTopic) {
-                return Arrays.asList(fetchTopicMetadata("test_topic"));
-            }
-            return null;
+            MetadataResponse.PartitionMetadata partitionMetadata = new 
MetadataResponse.PartitionMetadata(Errors.NONE, 1, null, new ArrayList<Node>(), 
new ArrayList<Node>());
+            MetadataResponse.TopicMetadata topicMetadata = new 
MetadataResponse.TopicMetadata(Errors.NONE, topic, true, 
Collections.singletonList(partitionMetadata));
+            return Collections.singleton(topicMetadata);
         }
     }
 }
\ No newline at end of file

http://git-wip-us.apache.org/repos/asf/kafka/blob/7837d3e5/streams/src/test/java/org/apache/kafka/test/MockInternalTopicManager.java
----------------------------------------------------------------------
diff --git 
a/streams/src/test/java/org/apache/kafka/test/MockInternalTopicManager.java 
b/streams/src/test/java/org/apache/kafka/test/MockInternalTopicManager.java
index af26880..28aca34 100644
--- a/streams/src/test/java/org/apache/kafka/test/MockInternalTopicManager.java
+++ b/streams/src/test/java/org/apache/kafka/test/MockInternalTopicManager.java
@@ -28,11 +28,12 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 public class MockInternalTopicManager extends InternalTopicManager {
 
     public Map<String, Integer> readyTopics = new HashMap<>();
-    public MockConsumer<byte[], byte[]> restoreConsumer;
+    private MockConsumer<byte[], byte[]> restoreConsumer;
 
     public MockInternalTopicManager(StreamsConfig streamsConfig, 
MockConsumer<byte[], byte[]> restoreConsumer) {
         super(new StreamsKafkaClient(streamsConfig), 0, 0);
@@ -41,15 +42,26 @@ public class MockInternalTopicManager extends 
InternalTopicManager {
     }
 
     @Override
-    public void makeReady(InternalTopicConfig topic, int numPartitions) {
-        readyTopics.put(topic.name(), numPartitions);
+    public void makeReady(final Map<InternalTopicConfig, Integer> topics) {
+        for (Map.Entry<InternalTopicConfig, Integer> entry : 
topics.entrySet()) {
+            readyTopics.put(entry.getKey().name(), entry.getValue());
 
-        List<PartitionInfo> partitions = new ArrayList<>();
-        for (int i = 0; i < numPartitions; i++) {
-            partitions.add(new PartitionInfo(topic.name(), i, null, null, 
null));
-        }
+            final List<PartitionInfo> partitions = new ArrayList<>();
+            for (int i = 0; i < entry.getValue(); i++) {
+                partitions.add(new PartitionInfo(entry.getKey().name(), i, 
null, null, null));
+            }
 
-        restoreConsumer.updatePartitions(topic.name(), partitions);
+            restoreConsumer.updatePartitions(entry.getKey().name(), 
partitions);
+        }
     }
 
+    @Override
+    public Map<String, Integer> getNumPartitions(final Set<String> topics) {
+        final Map<String, Integer> partitions = new HashMap<>();
+        for (String topic : topics) {
+            partitions.put(topic, restoreConsumer.partitionsFor(topic) == null 
?  null : restoreConsumer.partitionsFor(topic).size());
+        }
+
+        return partitions;
+    }
 }
\ No newline at end of file

Reply via email to