test: removing more infinite timeouts

test/integration/http2_integration_test.cc

@@ -1533,12 +1533,12 @@ void Http2FloodMitigationTest::beginSession() {
 
 Http2Frame Http2FloodMitigationTest::readFrame() {
   Http2Frame frame;
-  tcp_client_->waitForData(frame.HeaderSize);
+  EXPECT_TRUE(tcp_client_->waitForData(frame.HeaderSize));
   frame.setHeader(tcp_client_->data());
   tcp_client_->clearData(frame.HeaderSize);
   auto len = frame.payloadSize();
   if (len) {
-    tcp_client_->waitForData(len);
+    EXPECT_TRUE(tcp_client_->waitForData(len));
     frame.setPayload(tcp_client_->data());
     tcp_client_->clearData(len);
   }

test/integration/integration.cc

@@ -189,24 +189,28 @@ void IntegrationTcpClient::waitForData(const std::string& data, bool exact_match
   connection_->dispatcher().run(Event::Dispatcher::RunType::Block);
 }
 
-void IntegrationTcpClient::waitForData(size_t length) {
+AssertionResult IntegrationTcpClient::waitForData(size_t length,
+                                                  std::chrono::milliseconds timeout) {
   if (payload_reader_->data().size() >= length) {
-    return;
+    return AssertionSuccess();
   }
 
-  payload_reader_->setLengthToWaitFor(length);
-  connection_->dispatcher().run(Event::Dispatcher::RunType::Block);
+  return payload_reader_->waitForLength(length, timeout);
 }
 
 void IntegrationTcpClient::waitForDisconnect(bool ignore_spurious_events) {
+  Event::TimerPtr timeout_timer =
+      connection_->dispatcher().createTimer([this]() -> void { connection_->dispatcher().exit(); });
+  timeout_timer->enableTimer(TestUtility::DefaultTimeout);
+
   if (ignore_spurious_events) {
-    while (!disconnected_) {
+    while (!disconnected_ && timeout_timer->enabled()) {
       connection_->dispatcher().run(Event::Dispatcher::RunType::Block);
     }
   } else {
     connection_->dispatcher().run(Event::Dispatcher::RunType::Block);
-    EXPECT_TRUE(disconnected_);
   }
+  EXPECT_TRUE(disconnected_);
 }
 
 void IntegrationTcpClient::waitForHalfClose() {

test/integration/integration.h

@@ -104,7 +104,8 @@ class IntegrationTcpClient {
   void close();
   void waitForData(const std::string& data, bool exact_match = true);
   // wait for at least `length` bytes to be received
-  void waitForData(size_t length);
+  ABSL_MUST_USE_RESULT AssertionResult
+  waitForData(size_t length, std::chrono::milliseconds timeout = TestUtility::DefaultTimeout);
   void waitForDisconnect(bool ignore_spurious_events = false);
   void waitForHalfClose();
   void readDisable(bool disabled);

test/integration/tcp_proxy_integration_test.cc

@@ -70,13 +70,13 @@ TEST_P(TcpProxyIntegrationTest, TcpProxyUpstreamWritesFirst) {
   tcp_client->waitForData("ello", false);
 
   // Make sure length based wait works for the data already received
-  tcp_client->waitForData(5);
-  tcp_client->waitForData(4);
+  ASSERT_TRUE(tcp_client->waitForData(5));
+  ASSERT_TRUE(tcp_client->waitForData(4));
 
   // Drain part of the received message
   tcp_client->clearData(2);
   tcp_client->waitForData("llo");
-  tcp_client->waitForData(3);
+  ASSERT_TRUE(tcp_client->waitForData(3));
 
   ASSERT_TRUE(tcp_client->write("hello"));
   ASSERT_TRUE(fake_upstream_connection->waitForData(5));

test/integration/tcp_tunneling_integration_test.cc

@@ -302,7 +302,7 @@ TEST_P(TcpTunnelingIntegrationTest, Basic) {
 
   // Send data from upstream to downstream.
   upstream_request_->encodeData(12, false);
-  tcp_client->waitForData(12);
+  ASSERT_TRUE(tcp_client->waitForData(12));
 
   // Now send more data and close the TCP client. This should be treated as half close, so the data
   // should go through.
@@ -370,7 +370,7 @@ TEST_P(TcpTunnelingIntegrationTest, CloseUpstreamFirst) {
   // Send data from upstream to downstream with an end stream and make sure the data is received
   // before the connection is half-closed.
   upstream_request_->encodeData(12, true);
-  tcp_client->waitForData(12);
+  ASSERT_TRUE(tcp_client->waitForData(12));
   tcp_client->waitForHalfClose();
 
   // Attempt to send data upstream.

test/integration/utility.h

@@ -18,6 +18,8 @@
 #include "test/test_common/printers.h"
 #include "test/test_common/test_time.h"
 
+#include "gtest/gtest.h"
+
 namespace Envoy {
 /**
  * A buffering response decoder used for testing.
@@ -197,11 +199,29 @@ class WaitForPayloadReader : public Network::ReadFilterBaseImpl {
     data_to_wait_for_ = data;
     exact_match_ = exact_match;
   }
-  void setLengthToWaitFor(size_t length) {
+
+  ABSL_MUST_USE_RESULT testing::AssertionResult waitForLength(size_t length,
+                                                              std::chrono::milliseconds timeout) {
     ASSERT(!wait_for_length_);
     length_to_wait_for_ = length;
     wait_for_length_ = true;
+
+    Event::TimerPtr timeout_timer =
+        dispatcher_.createTimer([this]() -> void { dispatcher_.exit(); });
+    timeout_timer->enableTimer(timeout);
+
+    dispatcher_.run(Event::Dispatcher::RunType::Block);
+
+    if (timeout_timer->enabled()) {
+      timeout_timer->disableTimer();
+      return testing::AssertionSuccess();
+    }
+
+    length_to_wait_for_ = 0;
+    wait_for_length_ = false;
+    return testing::AssertionFailure() << "Timed out waiting for " << length << " bytes of data\n";
   }
+
   const std::string& data() { return data_; }
   bool readLastByte() { return read_end_stream_; }
   void clearData(size_t count = std::string::npos) { data_.erase(0, count); }