[VVP] Support any_of in categories decorator

[vvp/validation-scripts.git] / ice_validator / tests / helpers.py

diff --git a/ice_validator/tests/helpers.py b/ice_validator/tests/helpers.py
index 291464d..424dde1 100644 (file)
--- a/ice_validator/tests/helpers.py
+++ b/ice_validator/tests/helpers.py
@@ -124,16 +124,21 @@ def validates(*requirement_ids):
     return decorator
 
 
-def categories(*categories):
+def categories(*all_of, any_of=None):
+    any_of = set(any_of) if any_of else set()
+    all_of = set(all_of) if all_of else set()
+
     def decorator(func):
         @funcutils.wraps(func)
         def wrapper(*args, **kw):
             return func(*args, **kw)
 
-        wrapper.categories = categories
+        wrapper.all_categories = all_of
+        wrapper.any_categories = any_of
         return wrapper
 
-    decorator.categories = categories
+    decorator.all_categories = all_of
+    decorator.any_categories = any_of
     return decorator
 
 
@@ -393,7 +398,10 @@ def unzip(zip_path, target_dir):
     :param zip_path:    path to valid zip file
     :param target_dir:  directory to unzip zip_path
     """
-    check(zipfile.is_zipfile(zip_path), "{} is not a valid zipfile or does not exist".format(zip_path))
+    check(
+        zipfile.is_zipfile(zip_path),
+        "{} is not a valid zipfile or does not exist".format(zip_path),
+    )
     archive = zipfile.ZipFile(zip_path)
     if not os.path.exists(target_dir):
         os.makedirs(target_dir, exist_ok=True)