Loading...
tests/add_xzone_tests.py libmalloc-474.0.13 libmalloc-521.100.59
--- libmalloc/libmalloc-474.0.13/tests/add_xzone_tests.py
+++ libmalloc/libmalloc-521.100.59/tests/add_xzone_tests.py
@@ -3,8 +3,10 @@
 import copy
 import plistlib
 import sys
+import argparse
+from os.path import exists
 
-def add_xzone_tests(bats_plist_path):
+def add_xzone_tests(bats_plist_path, disable_xzone):
     with open(bats_plist_path, 'rb') as bats_plist_file:
         orig_bats_plist = plistlib.load(bats_plist_file)
 
@@ -13,11 +15,20 @@
 
     tests = []
     for test in orig_bats_plist['Tests']:
-        # Take the original test exactly as it was, unless it's xzone-only
+        # Keep the original test, unless it's xzone-only
         if 'Tags' not in test or 'xzone_only' not in test['Tags']:
-            tests.append(test)
+            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
+            tests.append(test_copy)
 
-        if 'Tags' in test and ('xzone' in test['Tags'] or 'xzone_only' in test['Tags']):
+        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)
 
@@ -39,16 +50,20 @@
 
                 xzone_test['CanonicalName'] = new_name
 
-            envvars = ['MallocSecureAllocator=1']
-            if 'perf' not in xzone_test['Tags']:
-                # This isn't a performance test, so we can use the debug variant
-                # of the library for extra assert coverage
-                envvars.append('DYLD_LIBRARY_PATH=/AppleInternal/Tests/libmalloc/assets')
-                envvars.append('DYLD_IMAGE_SUFFIX=_testdebug')
-
-                # However, using the debug variant does prevent leaks from
-                # working, so don't bother trying
-                envvars.append('DT_BYPASS_LEAKS_CHECK=1')
+            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)
@@ -63,7 +78,17 @@
         plistlib.dump(new_bats_plist, bats_plist_file, fmt=plistlib.FMT_BINARY)
 
 if __name__ == '__main__':
-    if len(sys.argv) != 2:
-        sys.exit('Expected bats plist as argument')
+    parser = argparse.ArgumentParser(description='Adds xzone malloc tests to BATS plist')
+    parser.add_argument('bats_plist')
+    parser.add_argument('-p', '--platform', help='Platform tests are being run on ()')
+    args = parser.parse_args()
 
-    add_xzone_tests(sys.argv[1])
+    assert exists(args.bats_plist), 'Expected bats plist as argument'
+
+    disable_xzone_tests = False
+    if args.platform == 'WatchOS':
+        # MallocSecureAllocator=1 is a no-op on arm64_32 (read: all watches),
+        # so don't generate .xzone test variants on watchOS
+        disable_xzone_tests = True
+
+    add_xzone_tests(args.bats_plist, disable_xzone_tests)