From 9fa3cb82bd54c74e829948e3ebe35ef93a0ed86a Mon Sep 17 00:00:00 2001
From: Derek Ostrander <derko@squareup.com>
Date: Fri, 9 Dec 2022 13:34:25 -0500
Subject: [PATCH] Add optional parameter of product_module_name This will be
 forwarded into the xctestrun file in the test name doesn't accurately reflect
 the module underneath

---
 xctestrunner/test_runner/ios_test_runner.py | 19 ++++++++++++++++---
 xctestrunner/test_runner/xctest_session.py  |  7 +++++--
 xctestrunner/test_runner/xctestrun.py       | 11 +++++++----
 3 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/xctestrunner/test_runner/ios_test_runner.py b/xctestrunner/test_runner/ios_test_runner.py
index c85800d..2759c26 100644
--- a/xctestrunner/test_runner/ios_test_runner.py
+++ b/xctestrunner/test_runner/ios_test_runner.py
@@ -117,7 +117,8 @@ def _Prepare(args):
           test_bundle=args.test_bundle_path,
           xctestrun_file_path=args.xctestrun,
           test_type=args.test_type,
-          signing_options=_GetJson(args.signing_options_json_path))
+          signing_options=_GetJson(args.signing_options_json_path),
+          product_module_name=args.product_module_name)
       session.SetLaunchOptions(_GetJson(args.launch_options_json_path))
 
   test_parser = subparsers.add_parser(
@@ -153,7 +154,8 @@ def _Test(args):
           test_bundle=args.test_bundle_path,
           xctestrun_file_path=args.xctestrun,
           test_type=args.test_type,
-          signing_options=_GetJson(args.signing_options_json_path))
+          signing_options=_GetJson(args.signing_options_json_path),
+          product_module_name=args.product_module_name)
       session.SetLaunchOptions(_GetJson(args.launch_options_json_path))
       return session.RunTest(args.id)
 
@@ -172,6 +174,11 @@ def _Test(args):
       help='The platform of the device. The value can be ios_device or '
            'ios_simulator.'
   )
+  optional_arguments.add_argument(
+      '--product_module_name',
+      help='The product module name that will be set in the xctestrun file.'
+  )
+
   test_parser.set_defaults(func=_Test)
 
 
@@ -197,7 +204,8 @@ def _RunSimulatorTest(args):
             test_bundle=args.test_bundle_path,
             xctestrun_file_path=args.xctestrun,
             test_type=args.test_type,
-            signing_options=_GetJson(args.signing_options_json_path))
+            signing_options=_GetJson(args.signing_options_json_path),
+          product_module_name=args.product_module_name)
         session.SetLaunchOptions(_GetJson(args.launch_options_json_path))
         if not hostless:
           try:
@@ -239,6 +247,11 @@ def _SimulatorTest(args):
            'prefix with simulator type and os version. '
            'E.g., New-iPhone 6 Plus-10.2.')
   test_parser.set_defaults(func=_SimulatorTest)
+  optional_arguments = test_parser.add_argument_group('Optional arguments')
+  optional_arguments.add_argument(
+      '--product_module_name',
+      help='The product module name that will be set in the xctestrun file.'
+  )
 
 
 def _BuildParser():
diff --git a/xctestrunner/test_runner/xctest_session.py b/xctestrunner/test_runner/xctest_session.py
index 1e879db..1e08e70 100644
--- a/xctestrunner/test_runner/xctest_session.py
+++ b/xctestrunner/test_runner/xctest_session.py
@@ -78,7 +78,8 @@ def __exit__(self, unused_type, unused_value, unused_traceback):
   # TODO(albertdai): Support bundle id as the value of app_under_test and
   # test_bundle.
   def Prepare(self, app_under_test=None, test_bundle=None,
-              xctestrun_file_path=None, test_type=None, signing_options=None):
+              xctestrun_file_path=None, test_type=None, 
+              signing_options=None, product_module_name=None):
     """Prepares the test session.
 
     If xctestrun_file is not provided, will use app under test and test bundle
@@ -94,6 +95,8 @@ def Prepare(self, app_under_test=None, test_bundle=None,
       test_type: ios_constants.TestType. The type of test bundle.
       signing_options: dict, the signing app options. See
           ios_constants.SIGNING_OPTIONS_JSON_HELP for details.
+      product_module_name: string, the name of the module that is being tested.
+          This will be forwarded into the xctestrun file.
 
     Raises:
       ios_errors.IllegalArgumentError:
@@ -141,7 +144,7 @@ def Prepare(self, app_under_test=None, test_bundle=None,
       if test_type != ios_constants.TestType.LOGIC_TEST:
         xctestrun_factory = xctestrun.XctestRunFactory(
             app_under_test_dir, test_bundle_dir, self._sdk, self._device_arch,
-            test_type, signing_options, self._work_dir)
+            test_type, signing_options, self._work_dir, product_module_name)
         self._xctestrun_obj = xctestrun_factory.GenerateXctestrun()
       else:
         self._logic_test_bundle = test_bundle_dir
diff --git a/xctestrunner/test_runner/xctestrun.py b/xctestrunner/test_runner/xctestrun.py
index f409fd6..553e32a 100644
--- a/xctestrunner/test_runner/xctestrun.py
+++ b/xctestrunner/test_runner/xctestrun.py
@@ -275,7 +275,8 @@ def __init__(self, app_under_test_dir, test_bundle_dir,
                sdk=ios_constants.SDK.IPHONESIMULATOR,
                device_arch=ios_constants.ARCH.X86_64,
                test_type=ios_constants.TestType.XCUITEST,
-               signing_options=None, work_dir=None):
+               signing_options=None, work_dir=None,
+               product_module_name=None):
     """Initializes the XctestRun object.
 
     If arg work_dir is provided, the original app under test file and test
@@ -292,6 +293,7 @@ def __init__(self, app_under_test_dir, test_bundle_dir,
       signing_options: dict, the signing app options. See
           ios_constants.SIGNING_OPTIONS_JSON_HELP for details.
       work_dir: string, work directory which contains run files.
+      product_module_name: string, forwarded into the xctestrun.
 
     Raises:
       IllegalArgumentError: when the sdk or test type is not supported.
@@ -302,6 +304,7 @@ def __init__(self, app_under_test_dir, test_bundle_dir,
     self._sdk = sdk
     self._device_arch = device_arch
     self._test_type = test_type
+    self._product_module_name = product_module_name
     if self._sdk == ios_constants.SDK.IPHONEOS:
       self._on_device = True
       self._signing_options = signing_options
@@ -511,7 +514,7 @@ def _GenerateTestRootForXcuitest(self):
         'DYLD_LIBRARY_PATH': '__TESTROOT__:%s/usr/lib' % developer_path
     }
     self._xctestrun_dict = {
-        'ProductModuleName': self._test_name.replace("-", "_"),
+        'ProductModuleName': self._product_module_name or self._test_name.replace("-", "_"),
         'IsUITestBundle': True,
         'SystemAttachmentLifetime': 'keepNever',
         'TestBundlePath': self._test_bundle_dir,
@@ -678,7 +681,7 @@ def _GenerateTestRootForXctest(self):
         'DYLD_LIBRARY_PATH': '__TESTROOT__:%s/usr/lib:' % developer_path
     }
     self._xctestrun_dict = {
-        'ProductModuleName': self._test_name.replace("-", "_"),
+        'ProductModuleName': self._product_module_name or self._test_name.replace("-", "_"),
         'TestHostPath': self._app_under_test_dir,
         'TestBundlePath': self._test_bundle_dir,
         'IsAppHostedTestBundle': True,
@@ -699,7 +702,7 @@ def _GenerateTestRootForLogicTest(self):
         'DYLD_LIBRARY_PATH': dyld_framework_path
     }
     self._xctestrun_dict = {
-        'ProductModuleName': self._test_name.replace("-", "_"),
+        'ProductModuleName': self._product_module_name or self._test_name.replace("-", "_"),
         'TestBundlePath': self._test_bundle_dir,
         'TestHostPath': xcode_info_util.GetXctestToolPath(self._sdk),
         'TestingEnvironmentVariables': test_envs,