From 36a909dda24e69922ef84a2f1e0c8c9009e5c0bc Mon Sep 17 00:00:00 2001
From: Marc Bonnici <marc.bonnici@arm.com>
Date: Fri, 31 Mar 2023 15:26:10 +0100
Subject: [PATCH] fw/ApkWorkload: Allow workloads to provide apk arguments

In some workloads it is beneficial to be able to provide arguments
when launching the required APK. Add a new `arpk_arguments` property
to allow a workload to provide a dict of parameter names and values
that should be used when launching the APK.

The python types of the parameters are used to determine the data
type provided to the APK. Currently supported types are string, bool,
int and float.
---
 wa/framework/workload.py | 18 ++++++++++--------
 wa/utils/android.py      | 27 +++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/wa/framework/workload.py b/wa/framework/workload.py
index 1b701d4e..645e2db8 100644
--- a/wa/framework/workload.py
+++ b/wa/framework/workload.py
@@ -23,7 +23,7 @@ except ImportError:
     from pipes import quote
 
 
-from wa.utils.android import get_cacheable_apk_info
+from wa.utils.android import get_cacheable_apk_info, build_apk_launch_command
 from wa.framework.plugin import TargetedPlugin, Parameter
 from wa.framework.resource import (ApkFile, ReventFile,
                                    File, loose_version_matching,
@@ -180,6 +180,7 @@ class ApkWorkload(Workload):
     activity = None
     view = None
     clear_data_on_reset = True
+    apk_arguments = {}
 
     # Set this to True to mark that this workload requires the target apk to be run
     # for initialisation purposes before the main run is performed.
@@ -289,7 +290,8 @@ class ApkWorkload(Workload):
                                   clear_data_on_reset=self.clear_data_on_reset,
                                   activity=self.activity,
                                   min_version=self.min_version,
-                                  max_version=self.max_version)
+                                  max_version=self.max_version,
+                                  apk_arguments=self.apk_arguments)
 
     def validate(self):
         if self.min_version and self.max_version:
@@ -686,7 +688,7 @@ class PackageHandler(object):
     def __init__(self, owner, install_timeout=300, version=None, variant=None,
                  package_name=None, strict=False, force_install=False, uninstall=False,
                  exact_abi=False, prefer_host_package=True, clear_data_on_reset=True,
-                 activity=None, min_version=None, max_version=None):
+                 activity=None, min_version=None, max_version=None, apk_arguments=None):
         self.logger = logging.getLogger('apk')
         self.owner = owner
         self.target = self.owner.target
@@ -709,6 +711,7 @@ class PackageHandler(object):
         self.apk_version = None
         self.logcat_log = None
         self.error_msg = None
+        self.apk_arguments = apk_arguments
 
     def initialize(self, context):
         self.resolve_package(context)
@@ -853,11 +856,10 @@ class PackageHandler(object):
         self.apk_version = host_version
 
     def start_activity(self):
-        if not self.activity:
-            cmd = 'am start -W {}'.format(self.apk_info.package)
-        else:
-            cmd = 'am start -W -n {}/{}'.format(self.apk_info.package,
-                                                self.activity)
+
+        cmd = build_apk_launch_command(self.apk_info.package, self.activity,
+                                       self.apk_arguments)
+
         output = self.target.execute(cmd)
         if 'Error:' in output:
             # this will dismiss any error dialogs
diff --git a/wa/utils/android.py b/wa/utils/android.py
index 7c70f3aa..176cea11 100644
--- a/wa/utils/android.py
+++ b/wa/utils/android.py
@@ -16,6 +16,7 @@
 import logging
 import os
 from datetime import datetime
+from shlex import quote
 
 from devlib.utils.android import ApkInfo as _ApkInfo
 
@@ -196,3 +197,29 @@ def get_cacheable_apk_info(path):
 
 
 apk_info_cache = ApkInfoCache()
+
+
+def build_apk_launch_command(package, activity=None, apk_args=None):
+    args_string = ''
+    if apk_args:
+        for k, v in apk_args.items():
+            if isinstance(v, str):
+                arg = '--es'
+                v = quote(v)
+            elif isinstance(v, float):
+                arg = '--ef'
+            elif isinstance(v, bool):
+                arg = '--ez'
+            elif isinstance(v, int):
+                arg = '--ei'
+            else:
+                raise ValueError('Unable to encode {} {}'.format(v, type(v)))
+
+            args_string = '{} {} {} {}'.format(args_string, arg, k, v)
+
+    if not activity:
+        cmd = 'am start -W {} {}'.format(package, args_string)
+    else:
+        cmd = 'am start -W -n {}/{} {}'.format(package, activity, args_string)
+
+    return cmd