Fix MigrationManagerTest flake (#591)

palantir · Dec 4, 2024 · 6aad1ab · 6aad1ab
1 parent b6d02f0
commit 6aad1ab
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/test/unit/org/apache/cassandra/service/MigrationManagerTest.java b/test/unit/org/apache/cassandra/service/MigrationManagerTest.java
@@ -20,8 +20,11 @@
 
 import java.net.InetAddress;
 import java.net.UnknownHostException;
+import java.util.Set;
 import java.util.UUID;
 
+import com.google.common.collect.ImmutableSet;
+import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 
@@ -39,6 +42,7 @@ public class MigrationManagerTest
     private static InetAddress HOST_2;
     private static InetAddress HOST_3;
     private static InetAddress HOST_4;
+    private static Set<InetAddress> hosts;
 
     @BeforeClass
     public static void setup() throws UnknownHostException
@@ -47,15 +51,20 @@ public static void setup() throws UnknownHostException
         HOST_2 = InetAddress.getByName("10.0.0.2");
         HOST_3 = InetAddress.getByName("10.0.0.3");
         HOST_4 = InetAddress.getByName("10.0.0.4");
+        hosts = ImmutableSet.of(HOST_1, HOST_2, HOST_3, HOST_4);
         MessagingService.instance().setVersion(HOST_1, MessagingService.VERSION_22);
         MessagingService.instance().setVersion(HOST_2, MessagingService.VERSION_22);
         MessagingService.instance().setVersion(HOST_3, MessagingService.VERSION_22);
         MessagingService.instance().setVersion(HOST_4, MessagingService.VERSION_22);
     }
 
+    @Before
+    public void before() {
+        hosts.forEach(host -> MigrationManager.removeEndpointFromSchemaPullVersion(UUID_1, host));
+    }
+
     @Test
-    public void shouldPullSchemaIfNoOutstandingRequests()
-    {
+    public void shouldPullSchemaIfNoOutstandingRequests() {
         assertTrue(MigrationManager.shouldPullSchemaFrom(HOST_1, UUID_1));
     }