From 1c32c738fb797a1ab22d222d49d5bc47fef30e7e Mon Sep 17 00:00:00 2001
From: Guogang Li <guogang@google.com>
Date: Wed, 30 Nov 2022 10:04:36 -0800
Subject: [PATCH] Call BWU callback after medium connection initiated

PiperOrigin-RevId: 491949519
---
 connections/implementation/base_pcp_handler.cc     | 14 +++++++++++---
 .../implementation/base_pcp_handler_test.cc        |  9 +++++----
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/connections/implementation/base_pcp_handler.cc b/connections/implementation/base_pcp_handler.cc
index 167b30dc..3b5081cf 100644
--- a/connections/implementation/base_pcp_handler.cc
+++ b/connections/implementation/base_pcp_handler.cc
@@ -17,6 +17,7 @@
 #include <algorithm>
 #include <cassert>
 #include <cinttypes>
+#include <cstddef>
 #include <cstdlib>
 #include <limits>
 #include <memory>
@@ -33,6 +34,7 @@
 #include "connections/connection_options.h"
 #include "connections/implementation/mediums/utils.h"
 #include "connections/implementation/offline_frames.h"
+#include "connections/medium_selector.h"
 #include "internal/platform/base64_utils.h"
 #include "internal/platform/bluetooth_utils.h"
 #include "internal/platform/logging.h"
@@ -1452,14 +1454,20 @@ void BasePcpHandler::EvaluateConnectionResult(ClientProxy* client,
     return;
   }
 
+  Medium medium =
+      channel_manager_->GetChannelForEndpoint(endpoint_id)->GetMedium();
   client->GetAnalyticsRecorder().OnConnectionEstablished(
-      endpoint_id,
-      channel_manager_->GetChannelForEndpoint(endpoint_id)->GetMedium(),
-      connection_info.connection_token);
+      endpoint_id, medium, connection_info.connection_token);
 
   // Invoke the client callback to let it know of the connection result.
   client->OnConnectionAccepted(endpoint_id);
 
+  // Report the current bandwidth to the client
+  client->OnBandwidthChanged(endpoint_id, medium);
+
+  NEARBY_LOGS(INFO) << "Connection accepted on Medium:"
+                    << proto::connections::Medium_Name(medium);
+
   // Kick off the bandwidth upgrade for incoming connections.
   if (connection_info.is_incoming &&
       AutoUpgradeBandwidth(client->GetAdvertisingOptions())) {
diff --git a/connections/implementation/base_pcp_handler_test.cc b/connections/implementation/base_pcp_handler_test.cc
index f8305094..18c55902 100644
--- a/connections/implementation/base_pcp_handler_test.cc
+++ b/connections/implementation/base_pcp_handler_test.cc
@@ -22,18 +22,18 @@
 #include "protobuf-matchers/protocol-buffer-matchers.h"
 #include "gtest/gtest.h"
 #include "absl/time/time.h"
-#include "connections/implementation/proto/offline_wire_formats.pb.h"
 #include "connections/implementation/base_endpoint_channel.h"
 #include "connections/implementation/bwu_manager.h"
 #include "connections/implementation/client_proxy.h"
 #include "connections/implementation/encryption_runner.h"
 #include "connections/implementation/offline_frames.h"
+#include "connections/implementation/proto/offline_wire_formats.pb.h"
 #include "connections/listeners.h"
 #include "connections/params.h"
 #include "internal/platform/byte_array.h"
+#include "internal/platform/count_down_latch.h"
 #include "internal/platform/exception.h"
 #include "internal/platform/medium_environment.h"
-#include "internal/platform/count_down_latch.h"
 #include "internal/platform/pipe.h"
 #include "proto/connections_enums.pb.h"
 
@@ -233,8 +233,8 @@ class BasePcpHandlerTest
         rejected_cb;
     StrictMock<MockFunction<void(const std::string& endpoint_id)>>
         disconnected_cb;
-    StrictMock<MockFunction<void(const std::string& endpoint_id,
-                                 std::int32_t quality)>>
+    StrictMock<
+        MockFunction<void(const std::string& endpoint_id, Medium medium)>>
         bandwidth_changed_cb;
   };
   struct MockDiscoveryListener {
@@ -647,6 +647,7 @@ TEST_P(BasePcpHandlerTest, OnIncomingFrameChangesState) {
   NEARBY_LOG(INFO, "Simulating remote accept: id=%s", endpoint_id.c_str());
   auto frame =
       parser::FromBytes(parser::ForConnectionResponse(Status::kSuccess));
+  EXPECT_CALL(mock_connection_listener_.bandwidth_changed_cb, Call).Times(1);
   pcp_handler.OnIncomingFrame(frame.result(), endpoint_id, &client,
                               connect_medium, packet_meta_data);
   NEARBY_LOGS(INFO) << "Closing connection: id=" << endpoint_id;