Loading...
tests/add_xzone_tests.py libmalloc-521.120.7 libmalloc-646.40.3
--- libmalloc/libmalloc-521.120.7/tests/add_xzone_tests.py
+++ libmalloc/libmalloc-646.40.3/tests/add_xzone_tests.py
@@ -1,10 +1,66 @@
 #!/usr/bin/env python3
 
+import argparse
 import copy
 import plistlib
-import sys
-import argparse
 from os.path import exists
+
+def create_xzone_test(test, nano_on_xzone):
+    xzone_test = copy.deepcopy(test)
+
+    extension = 'nano-on-xzone' if nano_on_xzone else 'xzone'
+
+    if 'TestName' in xzone_test:
+        orig_name = xzone_test['TestName']
+        xzone_test['TestName'] = orig_name + '.' + extension
+
+    if 'CanonicalName' in xzone_test:
+        # The CanonicalName for darwintests has a .<arch> suffix that we
+        # want to keep at the end, so insert our new component just
+        # before that
+        orig_name = xzone_test['CanonicalName']
+        components = orig_name.split('.')
+        if len(components) > 1:
+            components.insert(-1, extension)
+            new_name = '.'.join(components)
+        else:
+            new_name = orig_name + '.' + extension
+
+        xzone_test['CanonicalName'] = new_name
+
+    envvars = [
+        'MallocSecureAllocator=1',
+    ]
+
+    if nano_on_xzone:
+        envvars.append('MallocSecureAllocatorNano=0')
+        envvars.append('MallocNanoOnXzone=1')
+    else:
+        envvars.append('MallocSecureAllocatorNano=1')
+
+    if 'perf' not in xzone_test['Tags'] and \
+            'no_debug' not in xzone_test['Tags'] and \
+            not nano_on_xzone:
+        # This isn't a performance test or otherwise incompatible with
+        # the debug variant of the library, so we can use it for extra
+        # assert coverage.
+        envvars.append('DYLD_IMAGE_SUFFIX=_debug')
+
+    extend_env(xzone_test, envvars)
+
+    # xzone tests aren't eligible for VM testing because
+    # VMs report a single CPU cluster, which xzone malloc
+    # doesn't support
+    new_tags = []
+    for tag in test["Tags"]:
+        if tag == 'VM_PREFERRED':
+            new_tags.append("VM_NOT_PREFERRED")
+        else:
+            new_tags.append(tag)
+
+    xzone_test["Tags"] = new_tags
+
+    return xzone_test
 
 def add_xzone_tests(bats_plist_path, disable_xzone):
     with open(bats_plist_path, 'rb') as bats_plist_file:
@@ -15,67 +71,51 @@
 
     tests = []
     for test in orig_bats_plist['Tests']:
+        disable_pgm(test)
         # Keep the original test, unless it's xzone-only
         if 'Tags' not in test or 'xzone_only' not in test['Tags']:
             test_copy = copy.deepcopy(test)
             # Explicitly disable xzone malloc, so that the old allocator is
             # still tested on platforms that have xzone malloc by default
             envvars = ['MallocSecureAllocator=0']
-            if 'ShellEnv' in test_copy:
-                test_copy['ShellEnv'].extend(envvars)
-            else:
-                test_copy['ShellEnv'] = envvars
+            extend_env(test_copy, envvars)
             tests.append(test_copy)
 
-        if 'Tags' in test and ('xzone' in test['Tags'] or 'xzone_only' in test['Tags']) and \
-                not disable_xzone:
-            # This test has been tagged to run with xzone malloc
-            xzone_test = copy.deepcopy(test)
+        if 'Tags' in test and not disable_xzone:
+            if 'xzone' in test['Tags'] or 'xzone_only' in test['Tags']:
+                # This test has been tagged to run with xzone malloc
+                xzone_test = create_xzone_test(test, False)
+                tests.append(xzone_test)
 
-            if 'TestName' in xzone_test:
-                orig_name = xzone_test['TestName']
-                xzone_test['TestName'] = orig_name + '.xzone'
-
-            if 'CanonicalName' in xzone_test:
-                # The CanonicalName for darwintests has a .<arch> suffix that we
-                # want to keep at the end, so insert our new component just
-                # before that
-                orig_name = xzone_test['CanonicalName']
-                components = orig_name.split('.')
-                if len(components) > 1:
-                    components.insert(-1, 'xzone')
-                    new_name = '.'.join(components)
-                else:
-                    new_name = orig_name + '.xzone'
-
-                xzone_test['CanonicalName'] = new_name
-
-            envvars = [
-                'MallocSecureAllocator=1',
-                'MallocSecureAllocatorNano=1',
-            ]
-            if 'perf' not in xzone_test['Tags'] and \
-                    'no_debug' not in xzone_test['Tags'] and \
-                    not xzone_test.get('TestName', '').startswith('libmalloc.pgm'):
-                # This isn't a performance test or otherwise incompatible with
-                # the debug variant of the library, so we can use it for extra
-                # assert coverage.
-                #
-                # PGM tests can't use the debug variant until rdar://114584236
-                # is addressed
-                envvars.append('DYLD_IMAGE_SUFFIX=_debug')
-
-            if 'ShellEnv' in xzone_test:
-                xzone_test['ShellEnv'].extend(envvars)
-            else:
-                xzone_test['ShellEnv'] = envvars
-
-            tests.append(xzone_test)
+            if 'nano_on_xzone' in test['Tags']:
+                nano_on_xzone_test = create_xzone_test(test, True)
+                tests.append(nano_on_xzone_test)
 
     new_bats_plist['Tests'] = tests
 
     with open(bats_plist_path, 'wb') as bats_plist_file:
         plistlib.dump(new_bats_plist, bats_plist_file, fmt=plistlib.FMT_BINARY)
+
+
+# Note: The env is a list of `key=value`, not a dictionary
+def has_env_var(test, env_var_name):
+    env = test.get('ShellEnv', [])
+    key = env_var_name + '='  # exact match required
+    return any(v.startswith(key) for v in env)
+
+
+def extend_env(test, env_vars):
+    if 'ShellEnv' in test:
+        test['ShellEnv'].extend(env_vars)
+    else:
+        test['ShellEnv'] = env_vars
+
+
+# Disable PGM by setting `MallocProbGuard=0`, but respect existing uses
+def disable_pgm(test):
+    if not has_env_var(test, 'MallocProbGuard'):
+        extend_env(test, ['MallocProbGuard=0'])
+
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='Adds xzone malloc tests to BATS plist')