From 64d5ba426ecb604974cb035ba57f5dcf7e1c009e Mon Sep 17 00:00:00 2001
From: Bruno Cadonna <bruno@confluent.io>
Date: Thu, 3 Dec 2020 15:57:13 +0100
Subject: [PATCH] MINOR: Fix flaky test
 shouldQueryOnlyActivePartitionStoresByDefault

---
 .../kafka/streams/integration/StoreQueryIntegrationTest.java  | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/StoreQueryIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/StoreQueryIntegrationTest.java
index 712ae917b1859..9142314689758 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/StoreQueryIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/StoreQueryIntegrationTest.java
@@ -136,13 +136,13 @@ public void shouldQueryOnlyActivePartitionStoresByDefault() throws Exception {
 
             final boolean kafkaStreams1IsActive = (keyQueryMetadata.activeHost().port() % 2) == 1;
 
-            // Assert that only active is able to query for a key by default
-            assertThat(kafkaStreams1IsActive ? store1.get(key) : store2.get(key), is(notNullValue()));
             try {
                 if (kafkaStreams1IsActive) {
+                    assertThat(store1.get(key), is(notNullValue()));
                     assertThat(store2.get(key), is(nullValue()));
                 } else {
                     assertThat(store1.get(key), is(nullValue()));
+                    assertThat(store2.get(key), is(notNullValue()));
                 }
                 return true;
             } catch (final InvalidStateStoreException exception) {