avatar: Implement test decorator to enable an aconfig flag

Bug: 372300895
Change-Id: I2ebc8fda6e33acc93d7ce4a822b3f0bc47a9018f

avatar/__init__.py

@@ -42,6 +42,7 @@
 # public symbols
 __all__ = [
     'asynchronous',
+    'enableFlag',
     'parameterized',
     'rpc_except',
     'PandoraDevices',
@@ -206,6 +207,48 @@ def normalize(a: Any) -> Any:
     return wrapper
 
 
+def enableFlag(flag):
+    """ Enable aconfig flag.
+    Requires that the test class declares a devices: Optional[PandoraDevices] attribute.
+
+    Args:
+        flag: aconfig flag name including package, e.g.: 'com.android.bluetooth.flags.<flag_name>'
+
+    Raises:
+        AttributeError: when the 'devices' attribute is not found or not set
+        TypeError: when the provided flag argument is not a string
+    """
+
+    def decorator(func):
+
+        @functools.wraps(func)
+        def wrapper(self, *args, **kwargs):
+            devices = getattr(self, 'devices', None)
+
+            if not devices:
+                raise AttributeError(
+                    "Attribute 'devices' not found in test class or is None")
+
+            if not isinstance(devices, PandoraDevices):
+                raise TypeError(
+                    "devices attribute must be of a PandoraDevices type")
+
+            if not isinstance(flag, str):
+                raise TypeError("flag must be a string")
+
+            for server in devices._servers:
+                if isinstance(server, pandora_server.AndroidPandoraServer):
+                    server.device.adb.shell(
+                        ['device_config override bluetooth', flag,
+                         'true'])  # type: ignore
+                    break
+            return func(self, *args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 _T = TypeVar('_T')