From 6f5edc2615699abf5f4fc9f10e2e6d704572fd8d Mon Sep 17 00:00:00 2001
From: "justin.weng" <justin.weng@instaclustr.com>
Date: Thu, 17 Oct 2024 14:40:12 +1100
Subject: [PATCH] update tests

---
 shotover-proxy/tests/kafka_int_tests/mod.rs   | 13 ++++-------
 .../tests/kafka_int_tests/test_cases.rs       | 23 +++----------------
 2 files changed, 8 insertions(+), 28 deletions(-)

diff --git a/shotover-proxy/tests/kafka_int_tests/mod.rs b/shotover-proxy/tests/kafka_int_tests/mod.rs
index 45a40860d..e79381847 100644
--- a/shotover-proxy/tests/kafka_int_tests/mod.rs
+++ b/shotover-proxy/tests/kafka_int_tests/mod.rs
@@ -3,7 +3,6 @@ mod test_cases;
 use crate::shotover_process;
 use pretty_assertions::assert_eq;
 use rstest::rstest;
-use std::collections::VecDeque;
 use std::time::Duration;
 use std::time::Instant;
 use test_cases::produce_consume_partitions1;
@@ -393,10 +392,9 @@ async fn cluster_1_rack_multi_shotover_with_1_shotover_down(#[case] driver: Kafk
 
     // produce and consume messages, kill 1 shotover node and produce and consume more messages
     let connection_builder = KafkaConnectionBuilder::new(driver, "localhost:9192");
-    let mut shotover_nodes_to_kill = shotovers.drain(0..1).collect::<VecDeque<_>>();
+    let shotover_nodes_to_kill: Vec<_> = shotovers.drain(0..1).collect();
     test_cases::produce_consume_partitions1_shotover_nodes_go_down(
-        driver,
-        &mut shotover_nodes_to_kill,
+        shotover_nodes_to_kill,
         &connection_builder,
         "shotover_node_goes_down_test",
     )
@@ -453,10 +451,9 @@ async fn cluster_3_racks_multi_shotover_with_2_shotover_down(#[case] driver: Kaf
 
     // produce and consume messages, kill 2 shotover nodes and produce and consume more messages
     let connection_builder = KafkaConnectionBuilder::new(driver, "localhost:9193");
-    let mut shotover_nodes_to_kill = shotovers.drain(0..2).collect::<VecDeque<_>>();
+    let shotover_nodes_to_kill: Vec<_> = shotovers.drain(0..2).collect();
     test_cases::produce_consume_partitions1_shotover_nodes_go_down(
-        driver,
-        &mut shotover_nodes_to_kill,
+        shotover_nodes_to_kill,
         &connection_builder,
         "shotover_nodes_go_down_test",
     )
@@ -489,7 +486,7 @@ async fn cluster_3_racks_multi_shotover_with_2_shotover_down(#[case] driver: Kaf
 }
 
 #[rstest]
-// #[cfg_attr(feature = "kafka-cpp-driver-tests", case::cpp(KafkaDriver::Cpp))] //CPP driver may cause flaky tests.
+#[cfg_attr(feature = "kafka-cpp-driver-tests", case::cpp(KafkaDriver::Cpp))]
 #[case::java(KafkaDriver::Java)]
 #[tokio::test(flavor = "multi_thread")] // multi_thread is needed since java driver will block when consuming, causing shotover logs to not appear
 async fn cluster_3_racks_multi_shotover_with_1_shotover_missing(#[case] driver: KafkaDriver) {
diff --git a/shotover-proxy/tests/kafka_int_tests/test_cases.rs b/shotover-proxy/tests/kafka_int_tests/test_cases.rs
index 14f86fa75..9002c8f5c 100644
--- a/shotover-proxy/tests/kafka_int_tests/test_cases.rs
+++ b/shotover-proxy/tests/kafka_int_tests/test_cases.rs
@@ -1,5 +1,4 @@
 use futures::{stream::FuturesUnordered, StreamExt};
-use std::collections::VecDeque;
 use std::{collections::HashMap, time::Duration};
 use test_helpers::{
     connection::kafka::{
@@ -552,22 +551,10 @@ pub async fn produce_consume_partitions1_kafka_node_goes_down(
 }
 
 pub async fn produce_consume_partitions1_shotover_nodes_go_down(
-    driver: KafkaDriver,
-    shotover_nodes_to_kill: &mut VecDeque<BinProcess>,
+    shotover_nodes_to_kill: Vec<BinProcess>,
     connection_builder: &KafkaConnectionBuilder,
     topic_name: &str,
 ) {
-    if driver.is_cpp() {
-        // Skip this test for CPP driver.
-        // While the cpp driver has some retry capabilities,
-        // in many cases it will mark a shotover node as down for a single failed request
-        // and then immediately return the error to the caller, without waiting the full timeout period,
-        // since it has no more nodes to attempt sending to.
-        //
-        // So we skip this test on the CPP driver to avoid flaky tests.
-        return;
-    }
-
     {
         let admin = connection_builder.connect_admin().await;
         admin
@@ -608,15 +595,11 @@ pub async fn produce_consume_partitions1_shotover_nodes_go_down(
             })
             .await;
 
-        let num_nodes = shotover_nodes_to_kill.len();
         // kill shotover node(s)
-        for _ in 0..num_nodes {
+        for shotover_node in shotover_nodes_to_kill {
             tokio::time::timeout(
                 Duration::from_secs(10),
-                shotover_nodes_to_kill
-                    .pop_front()
-                    .unwrap()
-                    .shutdown_and_then_consume_events(&[]),
+                shotover_node.shutdown_and_then_consume_events(&[]),
             )
             .await
             .expect("Shotover did not shutdown within 10s");