Add at_least_B() decorator to restrict multi-device tests to Android B+ capable device.

Test: TH
Change-Id: Id5fde8c59560bb84bed062331121c3524f362c7b
diff --git a/staticlibs/testutils/host/python/apf_utils.py b/staticlibs/testutils/host/python/apf_utils.py
index 55ac860..8b9e155 100644
--- a/staticlibs/testutils/host/python/apf_utils.py
+++ b/staticlibs/testutils/host/python/apf_utils.py
@@ -18,6 +18,7 @@
 from mobly.controllers import android_device
 from mobly.controllers.android_device_lib.adb import AdbError
 from net_tests_utils.host.python import adb_utils, assert_utils
+import functools
 
 
 class PatternNotFoundException(Exception):
@@ -396,6 +397,20 @@
       f" {expected_version}",
   )
 
+def at_least_B():
+  def decorator(test_function):
+    @functools.wraps(test_function)
+    def wrapper(self, *args, **kwargs):
+      asserts.abort_class_if(
+        (not hasattr(self, 'client')) or (not hasattr(self.client, 'isAtLeastB')),
+        "client device is not B+"
+      )
+
+      asserts.abort_class_if(not self.client.isAtLeastB(), "not B+")
+      return test_function(self, *args, **kwargs)
+    return wrapper
+  return decorator
+
 class AdbOutputHandler:
   def __init__(self, ad, cmd):
     self._ad = ad
diff --git a/staticlibs/testutils/host/python/multi_devices_test_base.py b/staticlibs/testutils/host/python/multi_devices_test_base.py
index 677329a..72bac0c 100644
--- a/staticlibs/testutils/host/python/multi_devices_test_base.py
+++ b/staticlibs/testutils/host/python/multi_devices_test_base.py
@@ -53,3 +53,4 @@
         raise_on_exception=True,
     )
     self.client = self.clientDevice.connectivity_multi_devices_snippet
+    self.server = self.serverDevice.connectivity_multi_devices_snippet
diff --git a/tests/cts/multidevices/snippet/ConnectivityMultiDevicesSnippet.kt b/tests/cts/multidevices/snippet/ConnectivityMultiDevicesSnippet.kt
index 252052e..e1c6bf1 100644
--- a/tests/cts/multidevices/snippet/ConnectivityMultiDevicesSnippet.kt
+++ b/tests/cts/multidevices/snippet/ConnectivityMultiDevicesSnippet.kt
@@ -33,6 +33,8 @@
 import android.net.wifi.WifiInfo
 import android.net.wifi.WifiManager
 import android.net.wifi.WifiSsid
+import android.os.Build.VERSION.CODENAME
+import android.os.Build.VERSION.SDK_INT
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.compatibility.common.util.PropertyUtil
 import com.android.modules.utils.build.SdkLevel
@@ -62,6 +64,17 @@
         cbHelper.unregisterAll()
     }
 
+    private fun isAtLeastPreReleaseCodename(codeName: String): Boolean {
+        // Special case "REL", which means the build is not a pre-release build.
+        if ("REL".equals(CODENAME)) {
+            return false
+        }
+
+        // Otherwise lexically compare them. Return true if the build codename is equal to or
+        // greater than the requested codename.
+        return CODENAME.compareTo(codeName) >= 0
+    }
+
     @Rpc(description = "Check whether the device has wifi feature.")
     fun hasWifiFeature() = pm.hasSystemFeature(FEATURE_WIFI)
 
@@ -77,6 +90,11 @@
     @Rpc(description = "Return whether the Sdk level is at least V.")
     fun isAtLeastV() = SdkLevel.isAtLeastV()
 
+    @Rpc(description = "Check whether the device is at least B.")
+    fun isAtLeastB(): Boolean {
+        return SDK_INT >= 36 || (SDK_INT == 35 && isAtLeastPreReleaseCodename("Baklava"))
+    }
+
     @Rpc(description = "Return the API level that the VSR requirement must be fulfilled.")
     fun getVsrApiLevel() = PropertyUtil.getVsrApiLevel()