Fix: Allow both '*_test.py' and '*test_.py' test module naming conven…

…tion for nested tasks (#2155) * Allow both '*_test.py' and '*test_.py' test module naming convention for nested tasks Signed-off-by: Fabio Grätz <[email protected]> * Adapt docstring Signed-off-by: Fabio Grätz <[email protected]> --------- Signed-off-by: Fabio Grätz <[email protected]> Co-authored-by: Fabio Grätz <[email protected]>
flyteorg · Feb 5, 2024 · 0878397 · 0878397
1 parent f721eef
commit 0878397
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/flytekit/core/tracker.py b/flytekit/core/tracker.py
@@ -238,15 +238,17 @@ def foo_inner(*args, **kwargs):
 
 def istestfunction(func) -> bool:
     """
-    Returns true if the function is defined in a test module. A test module has to have `test_` as the prefix.
-    False in all other cases
+    Return true if the function is defined in a test module.
+
+    A test module has to have `test_` as the prefix or `_test` as the suffix.
+    False in all other cases.
     """
     mod = inspect.getmodule(func)
     if mod:
         mod_name = mod.__name__
         if "." in mod_name:
             mod_name = mod_name.split(".")[-1]
-        return mod_name.startswith("test_")
+        return mod_name.startswith("test_") or mod_name.endswith("_test")
     return False