Force BundleBasedDirectRunner in unit tests to bypass PrismRunner flakes

Author: vkarampudi

tfx/conftest.py

@@ -54,15 +54,16 @@ def custom_init(self, flags=None, **kwargs):
 
     has_other_runner = False
     for flag in flags_list:
-      if isinstance(flag, str) and flag.startswith('--runner=') and 'DirectRunner' not in flag:
+      if isinstance(flag, str) and flag.startswith('--runner='):
         has_other_runner = True
         break
 
     runner_kwarg = kwargs.get('runner')
-    if runner_kwarg and 'DirectRunner' not in str(runner_kwarg):
+    if runner_kwarg:
       has_other_runner = True
 
     if not has_other_runner:
+      flags_list.append('--runner=BundleBasedDirectRunner')
       if not any(isinstance(flag, str) and flag.startswith('--direct_running_mode=') for flag in flags_list):
         flags_list.append('--direct_running_mode=in_memory')