commit c8579866e275d980b4793b7e1ba5c96631490b07
author dschinazi <dschinazi@google.com> Wed Jul 24 18:20:20 2019 -0700
committer Copybara-Service <copybara-worker@google.com> Thu Jul 25 18:01:52 2019 -0700
tree 2eaa684f2df672bf0a4fddcc48b09b400940cadc
parent 76420db9e8095261c5666ca29571c1adecb1a3fa

Internal QUICHE change

PiperOrigin-RevId: 259859589
Change-Id: I33f879f5422c0fad5d3570dbb48b82bac23cefff

quic/test_tools/quic_test_client.cc

diff --git a/quic/test_tools/quic_test_client.cc b/quic/test_tools/quic_test_client.cc
index e2e33c7..ef92cb0 100644
--- a/quic/test_tools/quic_test_client.cc
+++ b/quic/test_tools/quic_test_client.cc
@@ -233,9 +233,14 @@
 }
 
 QuicConnectionId MockableQuicClient::GenerateNewConnectionId() {
-  return server_connection_id_overridden_
-             ? override_server_connection_id_
-             : QuicClient::GenerateNewConnectionId();
+  if (server_connection_id_overridden_) {
+    return override_server_connection_id_;
+  }
+  if (override_server_connection_id_length_ >= 0) {
+    return QuicUtils::CreateRandomConnectionId(
+        override_server_connection_id_length_);
+  }
+  return QuicClient::GenerateNewConnectionId();
 }
 
 void MockableQuicClient::UseConnectionId(
@@ -244,9 +249,20 @@
   override_server_connection_id_ = server_connection_id;
 }
 
+void MockableQuicClient::UseConnectionIdLength(
+    int server_connection_id_length) {
+  override_server_connection_id_length_ = server_connection_id_length;
+}
+
 QuicConnectionId MockableQuicClient::GetClientConnectionId() {
-  return client_connection_id_overridden_ ? override_client_connection_id_
-                                          : QuicClient::GetClientConnectionId();
+  if (client_connection_id_overridden_) {
+    return override_client_connection_id_;
+  }
+  if (override_client_connection_id_length_ >= 0) {
+    return QuicUtils::CreateRandomConnectionId(
+        override_client_connection_id_length_);
+  }
+  return QuicClient::GetClientConnectionId();
 }
 
 void MockableQuicClient::UseClientConnectionId(
@@ -255,6 +271,11 @@
   override_client_connection_id_ = client_connection_id;
 }
 
+void MockableQuicClient::UseClientConnectionIdLength(
+    int client_connection_id_length) {
+  override_client_connection_id_length_ = client_connection_id_length;
+}
+
 void MockableQuicClient::UseWriter(QuicPacketWriterWrapper* writer) {
   mockable_network_helper()->UseWriter(writer);
 }
@@ -769,12 +790,23 @@
   client_->UseConnectionId(server_connection_id);
 }
 
+void QuicTestClient::UseConnectionIdLength(int server_connection_id_length) {
+  DCHECK(!connected());
+  client_->UseConnectionIdLength(server_connection_id_length);
+}
+
 void QuicTestClient::UseClientConnectionId(
     QuicConnectionId client_connection_id) {
   DCHECK(!connected());
   client_->UseClientConnectionId(client_connection_id);
 }
 
+void QuicTestClient::UseClientConnectionIdLength(
+    int client_connection_id_length) {
+  DCHECK(!connected());
+  client_->UseClientConnectionIdLength(client_connection_id_length);
+}
+
 bool QuicTestClient::MigrateSocket(const QuicIpAddress& new_host) {
   return client_->MigrateSocket(new_host);
 }