Merge branch 'google3' to master, roll forward up to cl/355302206.

cpp/platform/public/BUILD

@@ -82,6 +82,7 @@ cc_library(
         "//platform/api:comm",
         "//platform/api:platform",
         "//platform/base",
+        "//platform/base:cancellation_flag",
         "//absl/container:flat_hash_map",
         "//absl/strings",
         "//webrtc/api:libjingle_peerconnection_api",

cpp/platform/public/ble.cc

@@ -127,13 +127,15 @@ bool BleMedium::StopAcceptingConnections(const std::string& service_id) {
 }
 
 BleSocket BleMedium::Connect(BlePeripheral& peripheral,
-                             const std::string& service_id) {
+                             const std::string& service_id,
+                             CancellationFlag* cancellation_flag) {
   {
     MutexLock lock(&mutex_);
     NEARBY_LOG(INFO, "BleMedium::Connect: peripheral=%p [impl=%p]", &peripheral,
                &peripheral.GetImpl());
   }
-  return BleSocket(impl_->Connect(peripheral.GetImpl(), service_id));
+  return BleSocket(
+      impl_->Connect(peripheral.GetImpl(), service_id, cancellation_flag));
 }
 
 }  // namespace nearby

cpp/platform/public/ble.h

@@ -18,6 +18,7 @@
 #include "platform/api/ble.h"
 #include "platform/api/platform.h"
 #include "platform/base/byte_array.h"
+#include "platform/base/cancellation_flag.h"
 #include "platform/base/input_stream.h"
 #include "platform/base/output_stream.h"
 #include "platform/public/bluetooth_adapter.h"
@@ -137,7 +138,8 @@ class BleMedium final {
 
   // Returns a new BleSocket. On Success, BleSocket::IsValid()
   // returns true.
-  BleSocket Connect(BlePeripheral& peripheral, const std::string& service_id);
+  BleSocket Connect(BlePeripheral& peripheral, const std::string& service_id,
+                    CancellationFlag* cancellation_flag);
 
   bool IsValid() const { return impl_ != nullptr; }
 

cpp/platform/public/ble_test.cc

@@ -26,12 +26,23 @@ namespace location {
 namespace nearby {
 namespace {
 
+using FeatureFlags = FeatureFlags::Flags;
+
+constexpr FeatureFlags kTestCases[] = {
+    FeatureFlags{
+        .enable_cancellation_flag = true,
+    },
+    FeatureFlags{
+        .enable_cancellation_flag = false,
+    },
+};
+
 constexpr absl::Duration kWaitDuration = absl::Milliseconds(1000);
 constexpr absl::string_view kServiceID{"com.google.location.nearby.apps.test"};
 constexpr absl::string_view kAdvertisementString{"\x0a\x0b\x0c\x0d"};
 constexpr absl::string_view kFastAdvertisementServiceUuid{"\xf3\xfe"};
 
-class BleMediumTest : public ::testing::Test {
+class BleMediumTest : public ::testing::TestWithParam<FeatureFlags> {
  protected:
   using DiscoveredPeripheralCallback = BleMedium::DiscoveredPeripheralCallback;
   using AcceptedConnectionCallback = BleMedium::AcceptedConnectionCallback;
@@ -41,6 +52,139 @@ class BleMediumTest : public ::testing::Test {
   MediumEnvironment& env_{MediumEnvironment::Instance()};
 };
 
+TEST_P(BleMediumTest, CanStartAcceptingConnectionsAndConnect) {
+  FeatureFlags feature_flags = GetParam();
+  env_.SetFeatureFlags(feature_flags);
+  env_.Start();
+  BluetoothAdapter adapter_a_;
+  BluetoothAdapter adapter_b_;
+  BleMedium ble_a{adapter_a_};
+  BleMedium ble_b{adapter_b_};
+  std::string service_id(kServiceID);
+  ByteArray advertisement_bytes{std::string(kAdvertisementString)};
+  std::string fast_advertisement_service_uuid(kFastAdvertisementServiceUuid);
+  CountDownLatch found_latch(1);
+  CountDownLatch accepted_latch(1);
+
+  BlePeripheral* discovered_peripheral = nullptr;
+  ble_a.StartScanning(
+      service_id, fast_advertisement_service_uuid,
+      DiscoveredPeripheralCallback{
+          .peripheral_discovered_cb =
+              [&found_latch, &discovered_peripheral](
+                  BlePeripheral& peripheral, const std::string& service_id,
+                  const ByteArray& advertisement_bytes,
+                  bool fast_advertisement) {
+                NEARBY_LOG(
+                    INFO,
+                    "Peripheral discovered: %s, %p, fast advertisement: %d",
+                    peripheral.GetName().c_str(), &peripheral,
+                    fast_advertisement);
+                discovered_peripheral = &peripheral;
+                found_latch.CountDown();
+              },
+      });
+  ble_b.StartAdvertising(service_id, advertisement_bytes,
+                         fast_advertisement_service_uuid);
+  ble_b.StartAcceptingConnections(
+      service_id,
+      AcceptedConnectionCallback{
+          .accepted_cb = [&accepted_latch](BleSocket socket,
+                                           const std::string& service_id) {
+            NEARBY_LOG(INFO, "Connection accepted: socket=%p, service_id=%s",
+                       &socket, service_id.c_str());
+            accepted_latch.CountDown();
+          }});
+  EXPECT_TRUE(found_latch.Await(kWaitDuration).result());
+
+  BleSocket socket_a;
+  EXPECT_FALSE(socket_a.IsValid());
+  {
+    SingleThreadExecutor client_executor;
+    client_executor.Execute(
+        [&ble_a, &socket_a, discovered_peripheral, &service_id]() {
+          CancellationFlag flag;
+          socket_a = ble_a.Connect(*discovered_peripheral, service_id, &flag);
+        });
+  }
+  EXPECT_TRUE(accepted_latch.Await(kWaitDuration).result());
+  EXPECT_TRUE(socket_a.IsValid());
+  ble_b.StopAdvertising(service_id);
+  ble_a.StopScanning(service_id);
+  env_.Stop();
+}
+
+TEST_P(BleMediumTest, CanCancelConnect) {
+  FeatureFlags feature_flags = GetParam();
+  env_.SetFeatureFlags(feature_flags);
+  env_.Start();
+  BluetoothAdapter adapter_a_;
+  BluetoothAdapter adapter_b_;
+  BleMedium ble_a{adapter_a_};
+  BleMedium ble_b{adapter_b_};
+  std::string service_id(kServiceID);
+  ByteArray advertisement_bytes{std::string(kAdvertisementString)};
+  std::string fast_advertisement_service_uuid(kFastAdvertisementServiceUuid);
+  CountDownLatch found_latch(1);
+  CountDownLatch accepted_latch(1);
+
+  BlePeripheral* discovered_peripheral = nullptr;
+  ble_a.StartScanning(
+      service_id, fast_advertisement_service_uuid,
+      DiscoveredPeripheralCallback{
+          .peripheral_discovered_cb =
+              [&found_latch, &discovered_peripheral](
+                  BlePeripheral& peripheral, const std::string& service_id,
+                  const ByteArray& advertisement_bytes,
+                  bool fast_advertisement) {
+                NEARBY_LOG(
+                    INFO,
+                    "Peripheral discovered: %s, %p, fast advertisement: %d",
+                    peripheral.GetName().c_str(), &peripheral,
+                    fast_advertisement);
+                discovered_peripheral = &peripheral;
+                found_latch.CountDown();
+              },
+      });
+  ble_b.StartAdvertising(service_id, advertisement_bytes,
+                         fast_advertisement_service_uuid);
+  ble_b.StartAcceptingConnections(
+      service_id,
+      AcceptedConnectionCallback{
+          .accepted_cb = [&accepted_latch](BleSocket socket,
+                                           const std::string& service_id) {
+            NEARBY_LOG(INFO, "Connection accepted: socket=%p, service_id=%s",
+                       &socket, service_id.c_str());
+            accepted_latch.CountDown();
+          }});
+  EXPECT_TRUE(found_latch.Await(kWaitDuration).result());
+
+  BleSocket socket_a;
+  EXPECT_FALSE(socket_a.IsValid());
+  {
+    SingleThreadExecutor client_executor;
+    client_executor.Execute(
+        [&ble_a, &socket_a, discovered_peripheral, &service_id]() {
+          CancellationFlag flag(true);
+          socket_a = ble_a.Connect(*discovered_peripheral, service_id, &flag);
+        });
+  }
+  // If FeatureFlag is disabled, Cancelled is false as no-op.
+  if (!feature_flags.enable_cancellation_flag) {
+    EXPECT_TRUE(accepted_latch.Await(kWaitDuration).result());
+    EXPECT_TRUE(socket_a.IsValid());
+  } else {
+    EXPECT_FALSE(accepted_latch.Await(kWaitDuration).result());
+    EXPECT_FALSE(socket_a.IsValid());
+  }
+  ble_b.StopAdvertising(service_id);
+  ble_a.StopScanning(service_id);
+  env_.Stop();
+}
+
+INSTANTIATE_TEST_SUITE_P(ParametrisedBleMediumTest, BleMediumTest,
+                         ::testing::ValuesIn(kTestCases));
+
 TEST_F(BleMediumTest, ConstructorDestructorWorks) {
   env_.Start();
   BluetoothAdapter adapter_a_;
@@ -156,65 +300,6 @@ TEST_F(BleMediumTest, CanStopDiscovery) {
   env_.Stop();
 }
 
-TEST_F(BleMediumTest, CanStartAcceptingConnectionsAndConnect) {
-  env_.Start();
-  BluetoothAdapter adapter_a_;
-  BluetoothAdapter adapter_b_;
-  BleMedium ble_a{adapter_a_};
-  BleMedium ble_b{adapter_b_};
-  std::string service_id(kServiceID);
-  ByteArray advertisement_bytes{std::string(kAdvertisementString)};
-  std::string fast_advertisement_service_uuid(kFastAdvertisementServiceUuid);
-  CountDownLatch found_latch(1);
-  CountDownLatch accepted_latch(1);
-
-  BlePeripheral* discovered_peripheral = nullptr;
-  ble_a.StartScanning(
-      service_id, fast_advertisement_service_uuid,
-      DiscoveredPeripheralCallback{
-          .peripheral_discovered_cb =
-              [&found_latch, &discovered_peripheral](
-                  BlePeripheral& peripheral, const std::string& service_id,
-                  const ByteArray& advertisement_bytes,
-                  bool fast_advertisement) {
-                NEARBY_LOG(
-                    INFO,
-                    "Peripheral discovered: %s, %p, fast advertisement: %d",
-                    peripheral.GetName().c_str(), &peripheral,
-                    fast_advertisement);
-                discovered_peripheral = &peripheral;
-                found_latch.CountDown();
-              },
-      });
-  ble_b.StartAdvertising(service_id, advertisement_bytes,
-                         fast_advertisement_service_uuid);
-  ble_b.StartAcceptingConnections(
-      service_id,
-      AcceptedConnectionCallback{
-          .accepted_cb = [&accepted_latch](BleSocket socket,
-                                           const std::string& service_id) {
-            NEARBY_LOG(INFO, "Connection accepted: socket=%p, service_id=%s",
-                       &socket, service_id.c_str());
-            accepted_latch.CountDown();
-          }});
-  EXPECT_TRUE(found_latch.Await(kWaitDuration).result());
-
-  BleSocket socket_a;
-  EXPECT_FALSE(socket_a.IsValid());
-  {
-    SingleThreadExecutor client_executor;
-    client_executor.Execute(
-        [&ble_a, &socket_a, discovered_peripheral, &service_id]() {
-          socket_a = ble_a.Connect(*discovered_peripheral, service_id);
-        });
-  }
-  EXPECT_TRUE(accepted_latch.Await(kWaitDuration).result());
-  EXPECT_TRUE(socket_a.IsValid());
-  ble_b.StopAdvertising(service_id);
-  ble_a.StopScanning(service_id);
-  env_.Stop();
-}
-
 }  // namespace
 }  // namespace nearby
 }  // namespace location

cpp/platform/public/bluetooth_classic.cc

@@ -23,12 +23,13 @@ namespace nearby {
 BluetoothClassicMedium::~BluetoothClassicMedium() { StopDiscovery(); }
 
 BluetoothSocket BluetoothClassicMedium::ConnectToService(
-    BluetoothDevice& remote_device, const std::string& service_uuid) {
+    BluetoothDevice& remote_device, const std::string& service_uuid,
+    CancellationFlag* cancellation_flag) {
   NEARBY_LOG(INFO,
              "BluetoothClassicMedium::ConnectToService: device=%p [impl=%p]",
              &remote_device, &remote_device.GetImpl());
-  return BluetoothSocket(
-      impl_->ConnectToService(remote_device.GetImpl(), service_uuid));
+  return BluetoothSocket(impl_->ConnectToService(
+      remote_device.GetImpl(), service_uuid, cancellation_flag));
 }
 
 bool BluetoothClassicMedium::StartDiscovery(DiscoveryCallback callback) {

cpp/platform/public/bluetooth_classic.h

@@ -21,6 +21,7 @@
 #include "platform/api/bluetooth_classic.h"
 #include "platform/api/platform.h"
 #include "platform/base/byte_array.h"
+#include "platform/base/cancellation_flag.h"
 #include "platform/base/exception.h"
 #include "platform/base/input_stream.h"
 #include "platform/base/listeners.h"
@@ -179,7 +180,8 @@ class BluetoothClassicMedium final {
   // Returns a new BluetoothSocket. On Success, BluetoothSocket::IsValid()
   // returns true.
   BluetoothSocket ConnectToService(BluetoothDevice& remote_device,
-                                   const std::string& service_uuid);
+                                   const std::string& service_uuid,
+                                   CancellationFlag* cancellation_flag);
 
   // https://developer.android.com/reference/android/bluetooth/BluetoothAdapter.html#listenUsingInsecureRfcommWithServiceRecord
   //

cpp/platform/public/bluetooth_classic_test.cc

@@ -29,7 +29,19 @@ namespace location {
 namespace nearby {
 namespace {
 
-class BluetoothClassicMediumTest : public ::testing::Test {
+using FeatureFlags = FeatureFlags::Flags;
+
+constexpr FeatureFlags kTestCases[] = {
+    FeatureFlags{
+        .enable_cancellation_flag = true,
+    },
+    FeatureFlags{
+        .enable_cancellation_flag = false,
+    },
+};
+
+class BluetoothClassicMediumTest
+    : public ::testing::TestWithParam<FeatureFlags> {
  protected:
   using DiscoveryCallback = BluetoothClassicMedium::DiscoveryCallback;
   BluetoothClassicMediumTest() {
@@ -66,6 +78,114 @@ class BluetoothClassicMediumTest : public ::testing::Test {
   std::unique_ptr<BluetoothClassicMedium> bt_b_;
 };
 
+TEST_P(BluetoothClassicMediumTest, CanConnectToService) {
+  FeatureFlags feature_flags = GetParam();
+  env_.SetFeatureFlags(feature_flags);
+
+  adapter_a_->SetScanMode(BluetoothAdapter::ScanMode::kConnectable);
+  CountDownLatch found_latch(1);
+  BluetoothDevice* discovered_device = nullptr;
+  bt_a_->StartDiscovery(DiscoveryCallback{
+      .device_discovered_cb =
+          [this, &found_latch, &discovered_device](BluetoothDevice& device) {
+            NEARBY_LOG(INFO, "Device discovered: %s", device.GetName().c_str());
+            EXPECT_EQ(device.GetName(), adapter_b_->GetName());
+            discovered_device = &device;
+            found_latch.CountDown();
+          },
+  });
+  adapter_b_->SetScanMode(BluetoothAdapter::ScanMode::kConnectableDiscoverable);
+  EXPECT_EQ(adapter_b_->GetScanMode(),
+            BluetoothAdapter::ScanMode::kConnectableDiscoverable);
+  EXPECT_TRUE(found_latch.Await(absl::Milliseconds(1000)).result());
+  std::string service_name{"service"};
+  std::string service_uuid("service-uuid");
+  BluetoothServerSocket server_socket =
+      bt_b_->ListenForService(service_name, service_uuid);
+  EXPECT_TRUE(server_socket.IsValid());
+  BluetoothSocket socket_a;
+  BluetoothSocket socket_b;
+  EXPECT_FALSE(socket_a.IsValid());
+  EXPECT_FALSE(socket_b.IsValid());
+  {
+    SingleThreadExecutor server_executor;
+    SingleThreadExecutor client_executor;
+    client_executor.Execute(
+        [this, &socket_a, discovered_device, &service_uuid, &server_socket]() {
+          CancellationFlag flag;
+          socket_a =
+              bt_a_->ConnectToService(*discovered_device, service_uuid, &flag);
+          if (!socket_a.IsValid()) server_socket.Close();
+        });
+    server_executor.Execute([&socket_b, &server_socket]() {
+      socket_b = server_socket.Accept();
+      if (!socket_b.IsValid()) server_socket.Close();
+    });
+  }
+  EXPECT_TRUE(socket_a.IsValid());
+  EXPECT_TRUE(socket_b.IsValid());
+  server_socket.Close();
+}
+
+TEST_P(BluetoothClassicMediumTest, CanCancelConnect) {
+  FeatureFlags feature_flags = GetParam();
+  env_.SetFeatureFlags(feature_flags);
+
+  adapter_a_->SetScanMode(BluetoothAdapter::ScanMode::kConnectable);
+  CountDownLatch found_latch(1);
+  BluetoothDevice* discovered_device = nullptr;
+  bt_a_->StartDiscovery(DiscoveryCallback{
+      .device_discovered_cb =
+          [this, &found_latch, &discovered_device](BluetoothDevice& device) {
+            NEARBY_LOG(INFO, "Device discovered: %s", device.GetName().c_str());
+            EXPECT_EQ(device.GetName(), adapter_b_->GetName());
+            discovered_device = &device;
+            found_latch.CountDown();
+          },
+  });
+  adapter_b_->SetScanMode(BluetoothAdapter::ScanMode::kConnectableDiscoverable);
+  EXPECT_EQ(adapter_b_->GetScanMode(),
+            BluetoothAdapter::ScanMode::kConnectableDiscoverable);
+  EXPECT_TRUE(found_latch.Await(absl::Milliseconds(1000)).result());
+  std::string service_name{"service"};
+  std::string service_uuid("service-uuid");
+  BluetoothServerSocket server_socket =
+      bt_b_->ListenForService(service_name, service_uuid);
+  EXPECT_TRUE(server_socket.IsValid());
+  BluetoothSocket socket_a;
+  BluetoothSocket socket_b;
+  EXPECT_FALSE(socket_a.IsValid());
+  EXPECT_FALSE(socket_b.IsValid());
+  {
+    SingleThreadExecutor server_executor;
+    SingleThreadExecutor client_executor;
+    client_executor.Execute(
+        [this, &socket_a, discovered_device, &service_uuid, &server_socket]() {
+          CancellationFlag flag(true);
+          socket_a =
+              bt_a_->ConnectToService(*discovered_device, service_uuid, &flag);
+          if (!socket_a.IsValid()) server_socket.Close();
+        });
+    server_executor.Execute([&socket_b, &server_socket]() {
+      socket_b = server_socket.Accept();
+      if (!socket_b.IsValid()) server_socket.Close();
+    });
+  }
+  // If FeatureFlag is disabled, Cancelled is false as no-op.
+  if (!feature_flags.enable_cancellation_flag) {
+    EXPECT_TRUE(socket_a.IsValid());
+    EXPECT_TRUE(socket_b.IsValid());
+  } else {
+    EXPECT_FALSE(socket_a.IsValid());
+    EXPECT_FALSE(socket_b.IsValid());
+  }
+  server_socket.Close();
+}
+
+INSTANTIATE_TEST_SUITE_P(ParametrisedBluetoothClassicMediumTest,
+                         BluetoothClassicMediumTest,
+                         ::testing::ValuesIn(kTestCases));
+
 TEST_F(BluetoothClassicMediumTest, ConstructorDestructorWorks) {
   // Make sure we can create functional adapters.
   ASSERT_TRUE(adapter_a_->IsValid());
@@ -163,51 +283,6 @@ TEST_F(BluetoothClassicMediumTest, CanListenForService) {
   server_socket.Close();
 }
 
-TEST_F(BluetoothClassicMediumTest, CanConnectToService) {
-  adapter_a_->SetScanMode(BluetoothAdapter::ScanMode::kConnectable);
-  CountDownLatch found_latch(1);
-  BluetoothDevice* discovered_device = nullptr;
-  bt_a_->StartDiscovery(DiscoveryCallback{
-      .device_discovered_cb =
-          [this, &found_latch, &discovered_device](BluetoothDevice& device) {
-            NEARBY_LOG(INFO, "Device discovered: %s", device.GetName().c_str());
-            EXPECT_EQ(device.GetName(), adapter_b_->GetName());
-            discovered_device = &device;
-            found_latch.CountDown();
-          },
-  });
-  adapter_b_->SetScanMode(BluetoothAdapter::ScanMode::kConnectableDiscoverable);
-  EXPECT_EQ(adapter_b_->GetScanMode(),
-            BluetoothAdapter::ScanMode::kConnectableDiscoverable);
-  EXPECT_TRUE(found_latch.Await(absl::Milliseconds(1000)).result());
-  std::string service_name{"service"};
-  std::string service_uuid("service-uuid");
-  BluetoothServerSocket server_socket =
-      bt_b_->ListenForService(service_name, service_uuid);
-  EXPECT_TRUE(server_socket.IsValid());
-  BluetoothSocket socket_a;
-  BluetoothSocket socket_b;
-  EXPECT_FALSE(socket_a.IsValid());
-  EXPECT_FALSE(socket_b.IsValid());
-  {
-    SingleThreadExecutor server_executor;
-    SingleThreadExecutor client_executor;
-    client_executor.Execute(
-        [this, &socket_a, discovered_device, &service_uuid, &server_socket]() {
-          socket_a = bt_a_->ConnectToService(*discovered_device, service_uuid);
-          if (!socket_a.IsValid()) server_socket.Close();
-        });
-    server_executor.Execute(
-        [&socket_b, &server_socket]() {
-          socket_b = server_socket.Accept();
-          if (!socket_b.IsValid()) server_socket.Close();
-        });
-  }
-  EXPECT_TRUE(socket_a.IsValid());
-  EXPECT_TRUE(socket_b.IsValid());
-  server_socket.Close();
-}
-
 }  // namespace
 }  // namespace nearby
 }  // namespace location

cpp/platform/public/wifi_lan.cc

@@ -136,13 +136,15 @@ bool WifiLanMedium::StopAcceptingConnections(const std::string& service_id) {
 }
 
 WifiLanSocket WifiLanMedium::Connect(WifiLanService& wifi_lan_service,
-                                     const std::string& service_id) {
+                                     const std::string& service_id,
+                                     CancellationFlag* cancellation_flag) {
   NEARBY_LOG(
       INFO,
       "WifiLanMedium::Connect: service=%p [impl=%p, service_info_name=%s]",
       &wifi_lan_service, &wifi_lan_service.GetImpl(),
       wifi_lan_service.GetServiceInfo().GetServiceInfoName().c_str());
-  return WifiLanSocket(impl_->Connect(wifi_lan_service.GetImpl(), service_id));
+  return WifiLanSocket(impl_->Connect(wifi_lan_service.GetImpl(), service_id,
+                                      cancellation_flag));
 }
 
 WifiLanService WifiLanMedium::GetRemoteService(const std::string& ip_address,

cpp/platform/public/wifi_lan.h

@@ -18,6 +18,7 @@
 #include "platform/api/platform.h"
 #include "platform/api/wifi_lan.h"
 #include "platform/base/byte_array.h"
+#include "platform/base/cancellation_flag.h"
 #include "platform/base/input_stream.h"
 #include "platform/base/nsd_service_info.h"
 #include "platform/base/output_stream.h"
@@ -151,7 +152,8 @@ class WifiLanMedium final {
   // Returns a new WifiLanSocket. On Success, WifiLanSocket::IsValid()
   // returns true.
   WifiLanSocket Connect(WifiLanService& wifi_lan_service,
-                        const std::string& service_id);
+                        const std::string& service_id,
+                        CancellationFlag* cancellation_flag);
 
   bool IsValid() const { return impl_ != nullptr; }
 

cpp/platform/public/wifi_lan_test.cc

@@ -27,12 +27,23 @@ namespace location {
 namespace nearby {
 namespace {
 
+using FeatureFlags = FeatureFlags::Flags;
+
+constexpr FeatureFlags kTestCases[] = {
+    FeatureFlags{
+        .enable_cancellation_flag = true,
+    },
+    FeatureFlags{
+        .enable_cancellation_flag = false,
+    },
+};
+
 constexpr absl::string_view kServiceID{"com.google.location.nearby.apps.test"};
 constexpr absl::string_view kServiceInfoName{"Simulated service info name"};
 constexpr absl::string_view kEndpointName{"Simulated endpoint name"};
 constexpr absl::string_view kEndpointInfoKey{"n"};
 
-class WifiLanMediumTest : public ::testing::Test {
+class WifiLanMediumTest : public ::testing::TestWithParam<FeatureFlags> {
  protected:
   using DiscoveredServiceCallback = WifiLanMedium::DiscoveredServiceCallback;
   using AcceptedConnectionCallback = WifiLanMedium::AcceptedConnectionCallback;
@@ -42,6 +53,142 @@ class WifiLanMediumTest : public ::testing::Test {
   MediumEnvironment& env_{MediumEnvironment::Instance()};
 };
 
+TEST_P(WifiLanMediumTest, CanStartAcceptingConnectionsAndConnect) {
+  FeatureFlags feature_flags = GetParam();
+  env_.SetFeatureFlags(feature_flags);
+  env_.Start();
+  WifiLanMedium wifi_a;
+  WifiLanMedium wifi_b;
+  std::string service_id(kServiceID);
+  std::string service_info_name{kServiceInfoName};
+  std::string endpoint_info_name{kEndpointName};
+  CountDownLatch found_latch(1);
+  CountDownLatch accepted_latch(1);
+
+  WifiLanService* discovered_service = nullptr;
+  wifi_a.StartDiscovery(
+      service_id,
+      DiscoveredServiceCallback{
+          .service_discovered_cb =
+              [&found_latch, &discovered_service](
+                  WifiLanService& service, const std::string& service_id) {
+                NEARBY_LOG(
+                    INFO, "Service discovered: %s, %p",
+                    service.GetServiceInfo().GetServiceInfoName().c_str(),
+                    &service);
+                discovered_service = &service;
+                found_latch.CountDown();
+              },
+      });
+
+  NsdServiceInfo nsd_service_info;
+  nsd_service_info.SetServiceInfoName(service_info_name);
+  nsd_service_info.SetTxtRecord(std::string(kEndpointInfoKey),
+                                endpoint_info_name);
+  wifi_b.StartAdvertising(service_id, nsd_service_info);
+  wifi_b.StartAcceptingConnections(
+      service_id,
+      AcceptedConnectionCallback{
+          .accepted_cb = [&accepted_latch](WifiLanSocket socket,
+                                           const std::string& service_id) {
+            NEARBY_LOG(INFO, "Connection accepted: socket=%p, service_id=%s",
+                       &socket, service_id.c_str());
+            accepted_latch.CountDown();
+          }});
+  EXPECT_TRUE(found_latch.Await(absl::Milliseconds(1000)).result());
+
+  WifiLanSocket socket_a;
+  EXPECT_FALSE(socket_a.IsValid());
+  {
+    SingleThreadExecutor client_executor;
+    client_executor.Execute(
+        [&wifi_a, &socket_a, discovered_service, &service_id]() {
+          CancellationFlag flag;
+          socket_a = wifi_a.Connect(*discovered_service, service_id, &flag);
+        });
+  }
+  EXPECT_TRUE(accepted_latch.Await(absl::Milliseconds(1000)).result());
+  EXPECT_TRUE(socket_a.IsValid());
+  wifi_b.StopAcceptingConnections(service_id);
+  wifi_b.StopAdvertising(service_id);
+  wifi_a.StopDiscovery(service_id);
+  env_.Stop();
+}
+
+TEST_P(WifiLanMediumTest, CanCancelConnect) {
+  FeatureFlags feature_flags = GetParam();
+  env_.SetFeatureFlags(feature_flags);
+  env_.Start();
+  WifiLanMedium wifi_a;
+  WifiLanMedium wifi_b;
+  std::string service_id(kServiceID);
+  std::string service_info_name{kServiceInfoName};
+  std::string endpoint_info_name{kEndpointName};
+  CountDownLatch found_latch(1);
+  CountDownLatch accepted_latch(1);
+
+  WifiLanService* discovered_service = nullptr;
+  wifi_a.StartDiscovery(
+      service_id,
+      DiscoveredServiceCallback{
+          .service_discovered_cb =
+              [&found_latch, &discovered_service](
+                  WifiLanService& service, const std::string& service_id) {
+                NEARBY_LOG(
+                    INFO, "Service discovered: %s, %p",
+                    service.GetServiceInfo().GetServiceInfoName().c_str(),
+                    &service);
+                discovered_service = &service;
+                found_latch.CountDown();
+              },
+      });
+
+  NsdServiceInfo nsd_service_info;
+  nsd_service_info.SetServiceInfoName(service_info_name);
+  nsd_service_info.SetTxtRecord(std::string(kEndpointInfoKey),
+                                endpoint_info_name);
+  wifi_b.StartAdvertising(service_id, nsd_service_info);
+  wifi_b.StartAcceptingConnections(
+      service_id,
+      AcceptedConnectionCallback{
+          .accepted_cb = [&accepted_latch](WifiLanSocket socket,
+                                           const std::string& service_id) {
+            NEARBY_LOG(INFO, "Connection accepted: socket=%p, service_id=%s",
+                       &socket, service_id.c_str());
+            accepted_latch.CountDown();
+          }});
+  EXPECT_TRUE(found_latch.Await(absl::Milliseconds(1000)).result());
+
+  WifiLanSocket socket_a;
+  EXPECT_FALSE(socket_a.IsValid());
+  {
+    SingleThreadExecutor client_executor;
+    client_executor.Execute(
+        [&wifi_a, &socket_a, discovered_service, &service_id]() {
+          // Make it as Cancelled.
+          CancellationFlag flag(true);
+          socket_a = wifi_a.Connect(*discovered_service, service_id, &flag);
+        });
+  }
+
+  // If FeatureFlag is disabled, Cancelled is false as no-op.
+  if (!feature_flags.enable_cancellation_flag) {
+    EXPECT_TRUE(accepted_latch.Await(absl::Milliseconds(1000)).result());
+    EXPECT_TRUE(socket_a.IsValid());
+  } else {
+    EXPECT_FALSE(accepted_latch.Await(absl::Milliseconds(1000)).result());
+    EXPECT_FALSE(socket_a.IsValid());
+  }
+
+  wifi_b.StopAcceptingConnections(service_id);
+  wifi_b.StopAdvertising(service_id);
+  wifi_a.StopDiscovery(service_id);
+  env_.Stop();
+}
+
+INSTANTIATE_TEST_SUITE_P(ParametrisedWifiLanMediumTest, WifiLanMediumTest,
+                         ::testing::ValuesIn(kTestCases));
+
 TEST_F(WifiLanMediumTest, ConstructorDestructorWorks) {
   env_.Start();
   WifiLanMedium wifi_a;
@@ -158,65 +305,6 @@ TEST_F(WifiLanMediumTest, CanStopDiscovery) {
   env_.Stop();
 }
 
-TEST_F(WifiLanMediumTest, CanStartAcceptingConnectionsAndConnect) {
-  env_.Start();
-  WifiLanMedium wifi_a;
-  WifiLanMedium wifi_b;
-  std::string service_id(kServiceID);
-  std::string service_info_name{kServiceInfoName};
-  std::string endpoint_info_name{kEndpointName};
-  CountDownLatch found_latch(1);
-  CountDownLatch accepted_latch(1);
-
-  WifiLanService* discovered_service = nullptr;
-  wifi_a.StartDiscovery(
-      service_id,
-      DiscoveredServiceCallback{
-          .service_discovered_cb =
-              [&found_latch, &discovered_service](
-                  WifiLanService& service, const std::string& service_id) {
-                NEARBY_LOG(
-                    INFO, "Service discovered: %s, %p",
-                    service.GetServiceInfo().GetServiceInfoName().c_str(),
-                    &service);
-                discovered_service = &service;
-                found_latch.CountDown();
-              },
-      });
-
-  NsdServiceInfo nsd_service_info;
-  nsd_service_info.SetServiceInfoName(service_info_name);
-  nsd_service_info.SetTxtRecord(std::string(kEndpointInfoKey),
-                                endpoint_info_name);
-  wifi_b.StartAdvertising(service_id, nsd_service_info);
-  wifi_b.StartAcceptingConnections(
-      service_id,
-      AcceptedConnectionCallback{
-          .accepted_cb = [&accepted_latch](WifiLanSocket socket,
-                                           const std::string& service_id) {
-            NEARBY_LOG(INFO, "Connection accepted: socket=%p, service_id=%s",
-                       &socket, service_id.c_str());
-            accepted_latch.CountDown();
-          }});
-  EXPECT_TRUE(found_latch.Await(absl::Milliseconds(1000)).result());
-
-  WifiLanSocket socket_a;
-  EXPECT_FALSE(socket_a.IsValid());
-  {
-    SingleThreadExecutor client_executor;
-    client_executor.Execute(
-        [&wifi_a, &socket_a, discovered_service, &service_id]() {
-          socket_a = wifi_a.Connect(*discovered_service, service_id);
-        });
-  }
-  EXPECT_TRUE(accepted_latch.Await(absl::Milliseconds(1000)).result());
-  EXPECT_TRUE(socket_a.IsValid());
-  wifi_b.StopAcceptingConnections(service_id);
-  wifi_b.StopAdvertising(service_id);
-  wifi_a.StopDiscovery(service_id);
-  env_.Stop();
-}
-
 }  // namespace
 }  // namespace nearby
 }  // namespace location