diff --git a/aikido_zen/middleware/set_rate_limit_group.py b/aikido_zen/middleware/set_rate_limit_group.py index 09f8f9c76..7fa30ac76 100644 --- a/aikido_zen/middleware/set_rate_limit_group.py +++ b/aikido_zen/middleware/set_rate_limit_group.py @@ -5,12 +5,12 @@ def set_rate_limit_group(group_id: Union[str, int]): - if not group_id: + if group_id is None or 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): + if isinstance(group_id, bool) or not isinstance(group_id, (str, int)): logger.warning("Group ID must be a string or a number") return group_id = str(group_id) diff --git a/aikido_zen/middleware/set_rate_limit_group_test.py b/aikido_zen/middleware/set_rate_limit_group_test.py index 019185871..2de57d246 100644 --- a/aikido_zen/middleware/set_rate_limit_group_test.py +++ b/aikido_zen/middleware/set_rate_limit_group_test.py @@ -39,6 +39,7 @@ def test_set_rate_limit_group_none_group_id(caplog): set_rate_limit_group(None) assert context1.rate_limit_group is None assert "Group ID cannot be empty." in caplog.text + assert "must be a string or a number" not in caplog.text def test_set_rate_limit_group_no_context(caplog): @@ -64,6 +65,38 @@ def test_set_rate_limit_group_non_string_group_id(caplog): assert context1.rate_limit_group == "123" +def test_set_rate_limit_group_zero_int_group_id(caplog): + context1 = test_utils.generate_and_set_context() + assert context1.rate_limit_group is None + set_rate_limit_group(0) + assert context1.rate_limit_group == "0" + assert "Group ID cannot be empty." not in caplog.text + + +def test_set_rate_limit_group_zero_string_group_id(caplog): + context1 = test_utils.generate_and_set_context() + assert context1.rate_limit_group is None + set_rate_limit_group("0") + assert context1.rate_limit_group == "0" + assert "Group ID cannot be empty." not in caplog.text + + +def test_set_rate_limit_group_false_group_id(caplog): + context1 = test_utils.generate_and_set_context() + assert context1.rate_limit_group is None + set_rate_limit_group(False) + 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_true_group_id(caplog): + context1 = test_utils.generate_and_set_context() + assert context1.rate_limit_group is None + set_rate_limit_group(True) + 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_non_string_group_id_non_number(caplog): context1 = test_utils.generate_and_set_context() assert context1.rate_limit_group is None