Allow numbers as group id

Author: bitterpanda63

aikido_zen/middleware/set_rate_limit_group.py

@@ -1,12 +1,20 @@
+from typing import Union
+
 from aikido_zen.context import get_current_context
 from aikido_zen.helpers.logging import logger
 
 
-def set_rate_limit_group(group_id: str):
-    if not group_id or not isinstance(group_id, str):
+def set_rate_limit_group(group_id: Union[str, int]):
+    if not group_id:
         logger.warning("Group ID cannot be empty.")
         return
 
+    # Check if it's string of number, ensure string.
+    if not isinstance(group_id, str) and not isinstance(group_id, int):
+        logger.warning("Group ID must be a string or a number")
+        return
+    group_id = str(group_id)
+
     context = get_current_context()
     if not context:
         logger.warning(

aikido_zen/middleware/set_rate_limit_group_test.py

@@ -84,8 +84,15 @@ def test_set_rate_limit_group_non_string_group_id(caplog):
     context1 = set_context_and_lifecycle()
     assert context1.rate_limit_group is None
     set_rate_limit_group(123)
+    assert context1.rate_limit_group == "123"
+
+
+def test_set_rate_limit_group_non_string_group_id_non_number(caplog):
+    context1 = set_context_and_lifecycle()
     assert context1.rate_limit_group is None
-    assert "Group ID cannot be empty." in caplog.text
+    set_rate_limit_group({"a": "b"})
+    assert context1.rate_limit_group is None
+    assert "Group ID must be a string or a number" in caplog.text
 
 
 def test_set_rate_limit_group_overwrite_existing_group():