mirror of
https://github.com/kuhyx/signal-bot.git
synced 2026-07-06 22:23:20 +02:00
Add abstract and generic rule system for command triggers
- Create rules.py module with Rule base class and implementations: - MinReactionsRule: Trigger if received n reactions or more - MaxUsageRule: Trigger if total usage did not exceed n - CooldownRule: Trigger if cooldown time has elapsed - NotUserRule: Trigger if last user was not person x - NotSameUserRule: Prevent same user from using consecutively - TimeWindowRule: Trigger if time is between t1 and t2 - AndRule/OrRule: Combine rules with AND/OR logic - GlobalCooldownRule: Rate limit across all commands - Update main.py to use the new RuleEngine - Add comprehensive test suite (60 tests) Co-authored-by: kuhyx <147418882+kuhyx@users.noreply.github.com>
This commit is contained in:
copilot-swe-agent[bot]
parent
7f3316d282
commit
8a6dbca7c9
71
main.py
71
main.py
@ -7,6 +7,21 @@ import json
|
||||
from datetime import datetime, time, timedelta
|
||||
from fastapi import FastAPI
|
||||
|
||||
from rules import (
|
||||
RuleEngine,
|
||||
CommandContext,
|
||||
GlobalCooldownRule,
|
||||
CooldownRule,
|
||||
MaxUsageRule,
|
||||
MinReactionsRule,
|
||||
NotUserRule,
|
||||
NotSameUserRule,
|
||||
TimeWindowRule,
|
||||
AndRule,
|
||||
OrRule,
|
||||
AlwaysPassRule,
|
||||
)
|
||||
|
||||
# Create FastAPI app
|
||||
app = FastAPI()
|
||||
|
||||
@ -17,7 +32,15 @@ SEND_URL = 'http://localhost:9922/v2/send'
|
||||
GROUP_ID = os.getenv('GROUP_ID', '')
|
||||
GROUP_ID_SEND = os.getenv('GROUP_ID_SEND', '')
|
||||
CAT_API = os.getenv('CAT_API', '')
|
||||
last_command_time = None
|
||||
|
||||
# Initialize the rule engine with a global 10-second cooldown (backward compatible)
|
||||
rule_engine = RuleEngine()
|
||||
rule_engine.add_global_rule(
|
||||
GlobalCooldownRule(
|
||||
cooldown=timedelta(seconds=10),
|
||||
failure_message="BEEP BOOP POCZEKAJ 10 SEKUND."
|
||||
)
|
||||
)
|
||||
warning_sent = False
|
||||
|
||||
class StringCounter:
|
||||
@ -154,23 +177,49 @@ async def scheduled_task(counter):
|
||||
send_message(counter.string_map, GROUP_ID_SEND)
|
||||
counter.string_map = {}
|
||||
|
||||
async def trigger_command(message_content, recipient):
|
||||
global last_command_time, warning_sent
|
||||
def extract_envelope_source_uuid(message):
|
||||
"""Extract sourceUuid from the raw message (for rule engine context)."""
|
||||
message_json = json.loads(message)
|
||||
return message_json.get('envelope', {}).get('sourceUuid', '')
|
||||
|
||||
|
||||
async def trigger_command(message_content, recipient, raw_message=None):
|
||||
global warning_sent
|
||||
message_value = message_message(message_content)
|
||||
|
||||
try:
|
||||
if message_value is not None and message_value[0] == "!":
|
||||
current_time = datetime.now()
|
||||
if last_command_time and current_time - last_command_time < timedelta(seconds=10):
|
||||
if not warning_sent:
|
||||
send_message("BEEP BOOP POCZEKAJ 10 SEKUND.", recipient)
|
||||
warning_sent = True
|
||||
return
|
||||
|
||||
|
||||
# Extract user ID for the rule engine
|
||||
user_id = ''
|
||||
if raw_message:
|
||||
user_id = extract_envelope_source_uuid(raw_message)
|
||||
|
||||
for command_triggers, command_function in command_map.items():
|
||||
if message_value in command_triggers:
|
||||
# Create command context for rule evaluation
|
||||
context = CommandContext(
|
||||
command=message_value,
|
||||
user_id=user_id,
|
||||
current_time=current_time,
|
||||
message_data=message_content
|
||||
)
|
||||
|
||||
# Evaluate rules
|
||||
should_trigger, failure_message = rule_engine.evaluate(context)
|
||||
|
||||
if not should_trigger:
|
||||
if failure_message and not warning_sent:
|
||||
send_message(failure_message, recipient)
|
||||
warning_sent = True
|
||||
return
|
||||
|
||||
# Execute the command
|
||||
await command_function(recipient)
|
||||
last_command_time = current_time
|
||||
|
||||
# Record usage for the rule engine
|
||||
rule_engine.record_usage(message_value, user_id, current_time)
|
||||
warning_sent = False
|
||||
break
|
||||
except TypeError:
|
||||
@ -181,7 +230,7 @@ async def trigger_command(message_content, recipient):
|
||||
async def send_to_group(message_content, counter, message):
|
||||
if message_group_id(message_content) == GROUP_ID:
|
||||
await count_messages(json.loads(message).get('envelope', {}), counter)
|
||||
await trigger_command(message_content, GROUP_ID_SEND)
|
||||
await trigger_command(message_content, GROUP_ID_SEND, message)
|
||||
|
||||
async def remove_attachment(attachment_id):
|
||||
response = requests.delete(REMOVE_ATTACHMENT_URL + attachment_id)
|
||||
|
||||
360
rules.py
Normal file
360
rules.py
Normal file
@ -0,0 +1,360 @@
|
||||
"""
|
||||
Abstract and generic rule system for triggering bot commands.
|
||||
|
||||
This module provides a flexible rule system that allows defining conditions
|
||||
under which commands should be triggered. Rules can be combined using
|
||||
logical AND/OR operations.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime, time, timedelta
|
||||
from typing import Optional, Dict, Any, List
|
||||
|
||||
|
||||
class CommandContext:
|
||||
"""Context object containing information about the current command execution."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
command: str,
|
||||
user_id: str,
|
||||
reaction_count: int = 0,
|
||||
current_time: Optional[datetime] = None,
|
||||
message_data: Optional[Dict[str, Any]] = None
|
||||
):
|
||||
self.command = command
|
||||
self.user_id = user_id
|
||||
self.reaction_count = reaction_count
|
||||
self.current_time = current_time or datetime.now()
|
||||
self.message_data = message_data or {}
|
||||
|
||||
|
||||
class CommandState:
|
||||
"""Tracks state for command usage (usage count, last user, last time)."""
|
||||
|
||||
def __init__(self):
|
||||
# Per-command state: command -> {usage_count, last_user, last_time}
|
||||
self._command_states: Dict[str, Dict[str, Any]] = {}
|
||||
|
||||
def get_state(self, command: str) -> Dict[str, Any]:
|
||||
"""Get state for a specific command."""
|
||||
if command not in self._command_states:
|
||||
self._command_states[command] = {
|
||||
'usage_count': 0,
|
||||
'last_user': None,
|
||||
'last_time': None
|
||||
}
|
||||
return self._command_states[command]
|
||||
|
||||
def record_usage(self, command: str, user_id: str, timestamp: Optional[datetime] = None):
|
||||
"""Record that a command was used."""
|
||||
state = self.get_state(command)
|
||||
state['usage_count'] += 1
|
||||
state['last_user'] = user_id
|
||||
state['last_time'] = timestamp or datetime.now()
|
||||
|
||||
def get_usage_count(self, command: str) -> int:
|
||||
"""Get total usage count for a command."""
|
||||
return self.get_state(command)['usage_count']
|
||||
|
||||
def get_last_user(self, command: str) -> Optional[str]:
|
||||
"""Get the last user who used a command."""
|
||||
return self.get_state(command)['last_user']
|
||||
|
||||
def get_last_time(self, command: str) -> Optional[datetime]:
|
||||
"""Get the last time a command was used."""
|
||||
return self.get_state(command)['last_time']
|
||||
|
||||
def reset_command(self, command: str):
|
||||
"""Reset state for a specific command."""
|
||||
if command in self._command_states:
|
||||
del self._command_states[command]
|
||||
|
||||
def reset_all(self):
|
||||
"""Reset all command states."""
|
||||
self._command_states.clear()
|
||||
|
||||
|
||||
class Rule(ABC):
|
||||
"""Abstract base class for command trigger rules."""
|
||||
|
||||
@abstractmethod
|
||||
def evaluate(self, context: CommandContext, state: CommandState) -> bool:
|
||||
"""
|
||||
Evaluate whether the rule passes.
|
||||
|
||||
Args:
|
||||
context: The current command context
|
||||
state: The command state tracker
|
||||
|
||||
Returns:
|
||||
True if the rule passes (command should be allowed), False otherwise
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_failure_message(self) -> Optional[str]:
|
||||
"""
|
||||
Get a message explaining why the rule failed.
|
||||
|
||||
Returns:
|
||||
A human-readable message or None if no message should be sent
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class MinReactionsRule(Rule):
|
||||
"""Rule that passes if the message has at least n reactions."""
|
||||
|
||||
def __init__(self, min_reactions: int, failure_message: Optional[str] = None):
|
||||
self.min_reactions = min_reactions
|
||||
self.failure_message = failure_message
|
||||
|
||||
def evaluate(self, context: CommandContext, state: CommandState) -> bool:
|
||||
return context.reaction_count >= self.min_reactions
|
||||
|
||||
def get_failure_message(self) -> Optional[str]:
|
||||
return self.failure_message
|
||||
|
||||
|
||||
class MaxUsageRule(Rule):
|
||||
"""Rule that passes if total command usage has not exceeded n."""
|
||||
|
||||
def __init__(self, max_usage: int, failure_message: Optional[str] = None):
|
||||
self.max_usage = max_usage
|
||||
self.failure_message = failure_message
|
||||
|
||||
def evaluate(self, context: CommandContext, state: CommandState) -> bool:
|
||||
current_usage = state.get_usage_count(context.command)
|
||||
return current_usage < self.max_usage
|
||||
|
||||
def get_failure_message(self) -> Optional[str]:
|
||||
return self.failure_message
|
||||
|
||||
|
||||
class CooldownRule(Rule):
|
||||
"""Rule that passes if sufficient time has passed since last command use."""
|
||||
|
||||
def __init__(self, cooldown: timedelta, failure_message: Optional[str] = None):
|
||||
self.cooldown = cooldown
|
||||
self.failure_message = failure_message
|
||||
|
||||
def evaluate(self, context: CommandContext, state: CommandState) -> bool:
|
||||
last_time = state.get_last_time(context.command)
|
||||
if last_time is None:
|
||||
return True
|
||||
return context.current_time - last_time >= self.cooldown
|
||||
|
||||
def get_failure_message(self) -> Optional[str]:
|
||||
return self.failure_message
|
||||
|
||||
|
||||
class NotUserRule(Rule):
|
||||
"""Rule that passes if the command was not last used by a specific user."""
|
||||
|
||||
def __init__(self, excluded_user: str, failure_message: Optional[str] = None):
|
||||
self.excluded_user = excluded_user
|
||||
self.failure_message = failure_message
|
||||
|
||||
def evaluate(self, context: CommandContext, state: CommandState) -> bool:
|
||||
last_user = state.get_last_user(context.command)
|
||||
if last_user is None:
|
||||
return True
|
||||
return last_user != self.excluded_user
|
||||
|
||||
def get_failure_message(self) -> Optional[str]:
|
||||
return self.failure_message
|
||||
|
||||
|
||||
class NotSameUserRule(Rule):
|
||||
"""Rule that passes if the current user is not the same as the last user."""
|
||||
|
||||
def __init__(self, failure_message: Optional[str] = None):
|
||||
self.failure_message = failure_message
|
||||
|
||||
def evaluate(self, context: CommandContext, state: CommandState) -> bool:
|
||||
last_user = state.get_last_user(context.command)
|
||||
if last_user is None:
|
||||
return True
|
||||
return last_user != context.user_id
|
||||
|
||||
def get_failure_message(self) -> Optional[str]:
|
||||
return self.failure_message
|
||||
|
||||
|
||||
class TimeWindowRule(Rule):
|
||||
"""Rule that passes if current time is within the specified window."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
start_time: time,
|
||||
end_time: time,
|
||||
failure_message: Optional[str] = None
|
||||
):
|
||||
self.start_time = start_time
|
||||
self.end_time = end_time
|
||||
self.failure_message = failure_message
|
||||
|
||||
def evaluate(self, context: CommandContext, state: CommandState) -> bool:
|
||||
current_time = context.current_time.time()
|
||||
|
||||
# Handle window that spans midnight (e.g., 22:00 to 06:00)
|
||||
if self.start_time <= self.end_time:
|
||||
return self.start_time <= current_time <= self.end_time
|
||||
else:
|
||||
return current_time >= self.start_time or current_time <= self.end_time
|
||||
|
||||
def get_failure_message(self) -> Optional[str]:
|
||||
return self.failure_message
|
||||
|
||||
|
||||
class AlwaysPassRule(Rule):
|
||||
"""Rule that always passes. Useful as a default or placeholder."""
|
||||
|
||||
def evaluate(self, context: CommandContext, state: CommandState) -> bool:
|
||||
return True
|
||||
|
||||
def get_failure_message(self) -> Optional[str]:
|
||||
return None
|
||||
|
||||
|
||||
class AlwaysFailRule(Rule):
|
||||
"""Rule that always fails. Useful for disabling commands."""
|
||||
|
||||
def __init__(self, failure_message: Optional[str] = None):
|
||||
self.failure_message = failure_message
|
||||
|
||||
def evaluate(self, context: CommandContext, state: CommandState) -> bool:
|
||||
return False
|
||||
|
||||
def get_failure_message(self) -> Optional[str]:
|
||||
return self.failure_message
|
||||
|
||||
|
||||
class AndRule(Rule):
|
||||
"""Composite rule that passes only if ALL child rules pass."""
|
||||
|
||||
def __init__(self, rules: List[Rule]):
|
||||
self.rules = rules
|
||||
self._failed_rule: Optional[Rule] = None
|
||||
|
||||
def evaluate(self, context: CommandContext, state: CommandState) -> bool:
|
||||
self._failed_rule = None
|
||||
for rule in self.rules:
|
||||
if not rule.evaluate(context, state):
|
||||
self._failed_rule = rule
|
||||
return False
|
||||
return True
|
||||
|
||||
def get_failure_message(self) -> Optional[str]:
|
||||
if self._failed_rule is not None:
|
||||
return self._failed_rule.get_failure_message()
|
||||
return None
|
||||
|
||||
|
||||
class OrRule(Rule):
|
||||
"""Composite rule that passes if ANY child rule passes."""
|
||||
|
||||
def __init__(self, rules: List[Rule], failure_message: Optional[str] = None):
|
||||
self.rules = rules
|
||||
self.failure_message = failure_message
|
||||
|
||||
def evaluate(self, context: CommandContext, state: CommandState) -> bool:
|
||||
for rule in self.rules:
|
||||
if rule.evaluate(context, state):
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_failure_message(self) -> Optional[str]:
|
||||
return self.failure_message
|
||||
|
||||
|
||||
class GlobalCooldownRule(Rule):
|
||||
"""Rule that enforces a cooldown across all commands (not per-command)."""
|
||||
|
||||
def __init__(self, cooldown: timedelta, failure_message: Optional[str] = None):
|
||||
self.cooldown = cooldown
|
||||
self.failure_message = failure_message
|
||||
self._last_global_time: Optional[datetime] = None
|
||||
|
||||
def evaluate(self, context: CommandContext, state: CommandState) -> bool:
|
||||
if self._last_global_time is None:
|
||||
return True
|
||||
return context.current_time - self._last_global_time >= self.cooldown
|
||||
|
||||
def record_usage(self, timestamp: Optional[datetime] = None):
|
||||
"""Record that a command was used globally."""
|
||||
self._last_global_time = timestamp or datetime.now()
|
||||
|
||||
def get_failure_message(self) -> Optional[str]:
|
||||
return self.failure_message
|
||||
|
||||
|
||||
class RuleEngine:
|
||||
"""
|
||||
Engine for evaluating rules and managing command state.
|
||||
|
||||
This provides a higher-level interface for working with rules.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.state = CommandState()
|
||||
self._command_rules: Dict[str, Rule] = {}
|
||||
self._default_rule: Rule = AlwaysPassRule()
|
||||
self._global_rules: List[Rule] = []
|
||||
|
||||
def set_command_rule(self, command: str, rule: Rule):
|
||||
"""Set the rule for a specific command."""
|
||||
self._command_rules[command] = rule
|
||||
|
||||
def set_command_rules(self, commands: tuple, rule: Rule):
|
||||
"""Set the same rule for multiple commands."""
|
||||
for command in commands:
|
||||
self._command_rules[command] = rule
|
||||
|
||||
def set_default_rule(self, rule: Rule):
|
||||
"""Set the default rule for commands without specific rules."""
|
||||
self._default_rule = rule
|
||||
|
||||
def add_global_rule(self, rule: Rule):
|
||||
"""Add a rule that applies to all commands."""
|
||||
self._global_rules.append(rule)
|
||||
|
||||
def clear_global_rules(self):
|
||||
"""Clear all global rules."""
|
||||
self._global_rules.clear()
|
||||
|
||||
def get_rule(self, command: str) -> Rule:
|
||||
"""Get the rule for a specific command."""
|
||||
return self._command_rules.get(command, self._default_rule)
|
||||
|
||||
def evaluate(self, context: CommandContext) -> tuple:
|
||||
"""
|
||||
Evaluate whether a command should be triggered.
|
||||
|
||||
Args:
|
||||
context: The command context
|
||||
|
||||
Returns:
|
||||
Tuple of (should_trigger: bool, failure_message: Optional[str])
|
||||
"""
|
||||
# First check global rules
|
||||
for global_rule in self._global_rules:
|
||||
if not global_rule.evaluate(context, self.state):
|
||||
return False, global_rule.get_failure_message()
|
||||
|
||||
# Then check command-specific rule
|
||||
rule = self.get_rule(context.command)
|
||||
if not rule.evaluate(context, self.state):
|
||||
return False, rule.get_failure_message()
|
||||
|
||||
return True, None
|
||||
|
||||
def record_usage(self, command: str, user_id: str, timestamp: Optional[datetime] = None):
|
||||
"""Record that a command was used."""
|
||||
self.state.record_usage(command, user_id, timestamp)
|
||||
|
||||
# Also record for global rules that need it
|
||||
for global_rule in self._global_rules:
|
||||
if isinstance(global_rule, GlobalCooldownRule):
|
||||
global_rule.record_usage(timestamp)
|
||||
732
test_rules.py
Normal file
732
test_rules.py
Normal file
@ -0,0 +1,732 @@
|
||||
"""
|
||||
Tests for the command trigger rules system.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from datetime import datetime, time, timedelta
|
||||
|
||||
from rules import (
|
||||
CommandContext,
|
||||
CommandState,
|
||||
RuleEngine,
|
||||
MinReactionsRule,
|
||||
MaxUsageRule,
|
||||
CooldownRule,
|
||||
NotUserRule,
|
||||
NotSameUserRule,
|
||||
TimeWindowRule,
|
||||
AlwaysPassRule,
|
||||
AlwaysFailRule,
|
||||
AndRule,
|
||||
OrRule,
|
||||
GlobalCooldownRule,
|
||||
)
|
||||
|
||||
|
||||
class TestCommandState:
|
||||
"""Tests for CommandState class."""
|
||||
|
||||
def test_initial_state(self):
|
||||
state = CommandState()
|
||||
assert state.get_usage_count("!test") == 0
|
||||
assert state.get_last_user("!test") is None
|
||||
assert state.get_last_time("!test") is None
|
||||
|
||||
def test_record_usage(self):
|
||||
state = CommandState()
|
||||
timestamp = datetime(2024, 1, 15, 12, 0, 0)
|
||||
state.record_usage("!test", "user123", timestamp)
|
||||
|
||||
assert state.get_usage_count("!test") == 1
|
||||
assert state.get_last_user("!test") == "user123"
|
||||
assert state.get_last_time("!test") == timestamp
|
||||
|
||||
def test_multiple_usages(self):
|
||||
state = CommandState()
|
||||
state.record_usage("!test", "user1")
|
||||
state.record_usage("!test", "user2")
|
||||
state.record_usage("!test", "user3")
|
||||
|
||||
assert state.get_usage_count("!test") == 3
|
||||
assert state.get_last_user("!test") == "user3"
|
||||
|
||||
def test_separate_commands(self):
|
||||
state = CommandState()
|
||||
state.record_usage("!cmd1", "user1")
|
||||
state.record_usage("!cmd2", "user2")
|
||||
|
||||
assert state.get_usage_count("!cmd1") == 1
|
||||
assert state.get_usage_count("!cmd2") == 1
|
||||
assert state.get_last_user("!cmd1") == "user1"
|
||||
assert state.get_last_user("!cmd2") == "user2"
|
||||
|
||||
def test_reset_command(self):
|
||||
state = CommandState()
|
||||
state.record_usage("!test", "user1")
|
||||
state.reset_command("!test")
|
||||
|
||||
assert state.get_usage_count("!test") == 0
|
||||
assert state.get_last_user("!test") is None
|
||||
|
||||
def test_reset_all(self):
|
||||
state = CommandState()
|
||||
state.record_usage("!cmd1", "user1")
|
||||
state.record_usage("!cmd2", "user2")
|
||||
state.reset_all()
|
||||
|
||||
assert state.get_usage_count("!cmd1") == 0
|
||||
assert state.get_usage_count("!cmd2") == 0
|
||||
|
||||
|
||||
class TestMinReactionsRule:
|
||||
"""Tests for MinReactionsRule."""
|
||||
|
||||
def test_passes_when_enough_reactions(self):
|
||||
rule = MinReactionsRule(min_reactions=5)
|
||||
context = CommandContext(command="!test", user_id="user1", reaction_count=5)
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_passes_when_more_reactions(self):
|
||||
rule = MinReactionsRule(min_reactions=5)
|
||||
context = CommandContext(command="!test", user_id="user1", reaction_count=10)
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_fails_when_fewer_reactions(self):
|
||||
rule = MinReactionsRule(min_reactions=5)
|
||||
context = CommandContext(command="!test", user_id="user1", reaction_count=4)
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is False
|
||||
|
||||
def test_failure_message(self):
|
||||
rule = MinReactionsRule(min_reactions=5, failure_message="Need more reactions!")
|
||||
assert rule.get_failure_message() == "Need more reactions!"
|
||||
|
||||
|
||||
class TestMaxUsageRule:
|
||||
"""Tests for MaxUsageRule."""
|
||||
|
||||
def test_passes_when_below_limit(self):
|
||||
rule = MaxUsageRule(max_usage=10)
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_passes_at_limit_minus_one(self):
|
||||
rule = MaxUsageRule(max_usage=10)
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
for _ in range(9):
|
||||
state.record_usage("!test", "user")
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_fails_at_limit(self):
|
||||
rule = MaxUsageRule(max_usage=10)
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
for _ in range(10):
|
||||
state.record_usage("!test", "user")
|
||||
|
||||
assert rule.evaluate(context, state) is False
|
||||
|
||||
def test_fails_above_limit(self):
|
||||
rule = MaxUsageRule(max_usage=10)
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
for _ in range(15):
|
||||
state.record_usage("!test", "user")
|
||||
|
||||
assert rule.evaluate(context, state) is False
|
||||
|
||||
|
||||
class TestCooldownRule:
|
||||
"""Tests for CooldownRule."""
|
||||
|
||||
def test_passes_when_no_previous_usage(self):
|
||||
rule = CooldownRule(cooldown=timedelta(seconds=10))
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_passes_when_cooldown_elapsed(self):
|
||||
rule = CooldownRule(cooldown=timedelta(seconds=10))
|
||||
state = CommandState()
|
||||
|
||||
past_time = datetime(2024, 1, 15, 12, 0, 0)
|
||||
current_time = datetime(2024, 1, 15, 12, 0, 15) # 15 seconds later
|
||||
|
||||
state.record_usage("!test", "user1", past_time)
|
||||
context = CommandContext(command="!test", user_id="user2", current_time=current_time)
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_fails_when_cooldown_not_elapsed(self):
|
||||
rule = CooldownRule(cooldown=timedelta(seconds=10))
|
||||
state = CommandState()
|
||||
|
||||
past_time = datetime(2024, 1, 15, 12, 0, 0)
|
||||
current_time = datetime(2024, 1, 15, 12, 0, 5) # 5 seconds later
|
||||
|
||||
state.record_usage("!test", "user1", past_time)
|
||||
context = CommandContext(command="!test", user_id="user2", current_time=current_time)
|
||||
|
||||
assert rule.evaluate(context, state) is False
|
||||
|
||||
def test_exact_cooldown_boundary(self):
|
||||
rule = CooldownRule(cooldown=timedelta(seconds=10))
|
||||
state = CommandState()
|
||||
|
||||
past_time = datetime(2024, 1, 15, 12, 0, 0)
|
||||
current_time = datetime(2024, 1, 15, 12, 0, 10) # Exactly 10 seconds later
|
||||
|
||||
state.record_usage("!test", "user1", past_time)
|
||||
context = CommandContext(command="!test", user_id="user2", current_time=current_time)
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
|
||||
class TestNotUserRule:
|
||||
"""Tests for NotUserRule."""
|
||||
|
||||
def test_passes_when_no_previous_usage(self):
|
||||
rule = NotUserRule(excluded_user="banned_user")
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_passes_when_different_user(self):
|
||||
rule = NotUserRule(excluded_user="banned_user")
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
state.record_usage("!test", "other_user")
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_fails_when_excluded_user(self):
|
||||
rule = NotUserRule(excluded_user="banned_user")
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
state.record_usage("!test", "banned_user")
|
||||
|
||||
assert rule.evaluate(context, state) is False
|
||||
|
||||
|
||||
class TestNotSameUserRule:
|
||||
"""Tests for NotSameUserRule."""
|
||||
|
||||
def test_passes_when_no_previous_usage(self):
|
||||
rule = NotSameUserRule()
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_passes_when_different_user(self):
|
||||
rule = NotSameUserRule()
|
||||
context = CommandContext(command="!test", user_id="user2")
|
||||
state = CommandState()
|
||||
state.record_usage("!test", "user1")
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_fails_when_same_user(self):
|
||||
rule = NotSameUserRule()
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
state.record_usage("!test", "user1")
|
||||
|
||||
assert rule.evaluate(context, state) is False
|
||||
|
||||
|
||||
class TestTimeWindowRule:
|
||||
"""Tests for TimeWindowRule."""
|
||||
|
||||
def test_passes_within_window(self):
|
||||
rule = TimeWindowRule(start_time=time(9, 0), end_time=time(17, 0))
|
||||
context = CommandContext(
|
||||
command="!test",
|
||||
user_id="user1",
|
||||
current_time=datetime(2024, 1, 15, 12, 0, 0) # 12:00
|
||||
)
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_passes_at_start(self):
|
||||
rule = TimeWindowRule(start_time=time(9, 0), end_time=time(17, 0))
|
||||
context = CommandContext(
|
||||
command="!test",
|
||||
user_id="user1",
|
||||
current_time=datetime(2024, 1, 15, 9, 0, 0) # 09:00
|
||||
)
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_passes_at_end(self):
|
||||
rule = TimeWindowRule(start_time=time(9, 0), end_time=time(17, 0))
|
||||
context = CommandContext(
|
||||
command="!test",
|
||||
user_id="user1",
|
||||
current_time=datetime(2024, 1, 15, 17, 0, 0) # 17:00
|
||||
)
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_fails_before_window(self):
|
||||
rule = TimeWindowRule(start_time=time(9, 0), end_time=time(17, 0))
|
||||
context = CommandContext(
|
||||
command="!test",
|
||||
user_id="user1",
|
||||
current_time=datetime(2024, 1, 15, 8, 0, 0) # 08:00
|
||||
)
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is False
|
||||
|
||||
def test_fails_after_window(self):
|
||||
rule = TimeWindowRule(start_time=time(9, 0), end_time=time(17, 0))
|
||||
context = CommandContext(
|
||||
command="!test",
|
||||
user_id="user1",
|
||||
current_time=datetime(2024, 1, 15, 18, 0, 0) # 18:00
|
||||
)
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is False
|
||||
|
||||
def test_overnight_window_passes_evening(self):
|
||||
# Window from 22:00 to 06:00 (overnight)
|
||||
rule = TimeWindowRule(start_time=time(22, 0), end_time=time(6, 0))
|
||||
context = CommandContext(
|
||||
command="!test",
|
||||
user_id="user1",
|
||||
current_time=datetime(2024, 1, 15, 23, 0, 0) # 23:00
|
||||
)
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_overnight_window_passes_morning(self):
|
||||
# Window from 22:00 to 06:00 (overnight)
|
||||
rule = TimeWindowRule(start_time=time(22, 0), end_time=time(6, 0))
|
||||
context = CommandContext(
|
||||
command="!test",
|
||||
user_id="user1",
|
||||
current_time=datetime(2024, 1, 15, 4, 0, 0) # 04:00
|
||||
)
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_overnight_window_fails_midday(self):
|
||||
# Window from 22:00 to 06:00 (overnight)
|
||||
rule = TimeWindowRule(start_time=time(22, 0), end_time=time(6, 0))
|
||||
context = CommandContext(
|
||||
command="!test",
|
||||
user_id="user1",
|
||||
current_time=datetime(2024, 1, 15, 12, 0, 0) # 12:00
|
||||
)
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is False
|
||||
|
||||
|
||||
class TestAlwaysPassRule:
|
||||
"""Tests for AlwaysPassRule."""
|
||||
|
||||
def test_always_passes(self):
|
||||
rule = AlwaysPassRule()
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_no_failure_message(self):
|
||||
rule = AlwaysPassRule()
|
||||
assert rule.get_failure_message() is None
|
||||
|
||||
|
||||
class TestAlwaysFailRule:
|
||||
"""Tests for AlwaysFailRule."""
|
||||
|
||||
def test_always_fails(self):
|
||||
rule = AlwaysFailRule()
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is False
|
||||
|
||||
def test_failure_message(self):
|
||||
rule = AlwaysFailRule(failure_message="Command disabled")
|
||||
assert rule.get_failure_message() == "Command disabled"
|
||||
|
||||
|
||||
class TestAndRule:
|
||||
"""Tests for AndRule composite rule."""
|
||||
|
||||
def test_passes_when_all_pass(self):
|
||||
rule = AndRule([
|
||||
AlwaysPassRule(),
|
||||
AlwaysPassRule(),
|
||||
AlwaysPassRule()
|
||||
])
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_fails_when_one_fails(self):
|
||||
rule = AndRule([
|
||||
AlwaysPassRule(),
|
||||
AlwaysFailRule(),
|
||||
AlwaysPassRule()
|
||||
])
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is False
|
||||
|
||||
def test_failure_message_from_failed_rule(self):
|
||||
rule = AndRule([
|
||||
AlwaysPassRule(),
|
||||
AlwaysFailRule(failure_message="This one failed"),
|
||||
AlwaysPassRule()
|
||||
])
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
rule.evaluate(context, state)
|
||||
assert rule.get_failure_message() == "This one failed"
|
||||
|
||||
def test_empty_rules_pass(self):
|
||||
rule = AndRule([])
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
|
||||
class TestOrRule:
|
||||
"""Tests for OrRule composite rule."""
|
||||
|
||||
def test_passes_when_one_passes(self):
|
||||
rule = OrRule([
|
||||
AlwaysFailRule(),
|
||||
AlwaysPassRule(),
|
||||
AlwaysFailRule()
|
||||
])
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_passes_when_all_pass(self):
|
||||
rule = OrRule([
|
||||
AlwaysPassRule(),
|
||||
AlwaysPassRule(),
|
||||
AlwaysPassRule()
|
||||
])
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_fails_when_all_fail(self):
|
||||
rule = OrRule([
|
||||
AlwaysFailRule(),
|
||||
AlwaysFailRule(),
|
||||
AlwaysFailRule()
|
||||
])
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is False
|
||||
|
||||
def test_failure_message(self):
|
||||
rule = OrRule([
|
||||
AlwaysFailRule(),
|
||||
AlwaysFailRule()
|
||||
], failure_message="None passed")
|
||||
|
||||
assert rule.get_failure_message() == "None passed"
|
||||
|
||||
def test_empty_rules_fail(self):
|
||||
rule = OrRule([])
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is False
|
||||
|
||||
|
||||
class TestGlobalCooldownRule:
|
||||
"""Tests for GlobalCooldownRule."""
|
||||
|
||||
def test_passes_when_no_previous_usage(self):
|
||||
rule = GlobalCooldownRule(cooldown=timedelta(seconds=10))
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
state = CommandState()
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_passes_when_cooldown_elapsed(self):
|
||||
rule = GlobalCooldownRule(cooldown=timedelta(seconds=10))
|
||||
state = CommandState()
|
||||
|
||||
past_time = datetime(2024, 1, 15, 12, 0, 0)
|
||||
current_time = datetime(2024, 1, 15, 12, 0, 15) # 15 seconds later
|
||||
|
||||
rule.record_usage(past_time)
|
||||
context = CommandContext(command="!test", user_id="user1", current_time=current_time)
|
||||
|
||||
assert rule.evaluate(context, state) is True
|
||||
|
||||
def test_fails_when_cooldown_not_elapsed(self):
|
||||
rule = GlobalCooldownRule(cooldown=timedelta(seconds=10))
|
||||
state = CommandState()
|
||||
|
||||
past_time = datetime(2024, 1, 15, 12, 0, 0)
|
||||
current_time = datetime(2024, 1, 15, 12, 0, 5) # 5 seconds later
|
||||
|
||||
rule.record_usage(past_time)
|
||||
context = CommandContext(command="!test", user_id="user1", current_time=current_time)
|
||||
|
||||
assert rule.evaluate(context, state) is False
|
||||
|
||||
|
||||
class TestRuleEngine:
|
||||
"""Tests for RuleEngine class."""
|
||||
|
||||
def test_default_always_pass(self):
|
||||
engine = RuleEngine()
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
|
||||
should_trigger, message = engine.evaluate(context)
|
||||
assert should_trigger is True
|
||||
assert message is None
|
||||
|
||||
def test_command_specific_rule(self):
|
||||
engine = RuleEngine()
|
||||
engine.set_command_rule("!test", AlwaysFailRule(failure_message="Nope"))
|
||||
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
should_trigger, message = engine.evaluate(context)
|
||||
|
||||
assert should_trigger is False
|
||||
assert message == "Nope"
|
||||
|
||||
def test_command_specific_rule_different_command(self):
|
||||
engine = RuleEngine()
|
||||
engine.set_command_rule("!test", AlwaysFailRule(failure_message="Nope"))
|
||||
|
||||
# Different command should use default rule
|
||||
context = CommandContext(command="!other", user_id="user1")
|
||||
should_trigger, message = engine.evaluate(context)
|
||||
|
||||
assert should_trigger is True
|
||||
assert message is None
|
||||
|
||||
def test_set_command_rules_multiple(self):
|
||||
engine = RuleEngine()
|
||||
engine.set_command_rules(("!cat", "!kot", "!meow"), AlwaysFailRule())
|
||||
|
||||
for cmd in ["!cat", "!kot", "!meow"]:
|
||||
context = CommandContext(command=cmd, user_id="user1")
|
||||
should_trigger, _ = engine.evaluate(context)
|
||||
assert should_trigger is False
|
||||
|
||||
def test_global_rule_blocks(self):
|
||||
engine = RuleEngine()
|
||||
engine.add_global_rule(AlwaysFailRule(failure_message="Global block"))
|
||||
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
should_trigger, message = engine.evaluate(context)
|
||||
|
||||
assert should_trigger is False
|
||||
assert message == "Global block"
|
||||
|
||||
def test_global_rule_evaluated_first(self):
|
||||
engine = RuleEngine()
|
||||
engine.add_global_rule(AlwaysFailRule(failure_message="Global"))
|
||||
engine.set_command_rule("!test", AlwaysFailRule(failure_message="Command"))
|
||||
|
||||
context = CommandContext(command="!test", user_id="user1")
|
||||
_, message = engine.evaluate(context)
|
||||
|
||||
# Global rule should be checked first
|
||||
assert message == "Global"
|
||||
|
||||
def test_record_usage(self):
|
||||
engine = RuleEngine()
|
||||
engine.record_usage("!test", "user1")
|
||||
|
||||
assert engine.state.get_usage_count("!test") == 1
|
||||
assert engine.state.get_last_user("!test") == "user1"
|
||||
|
||||
def test_complex_rule_combination(self):
|
||||
"""Test combining multiple rules."""
|
||||
engine = RuleEngine()
|
||||
|
||||
# Command can only be used 5 times total, with a 30-second cooldown,
|
||||
# and only between 9am and 5pm
|
||||
rule = AndRule([
|
||||
MaxUsageRule(max_usage=5, failure_message="Max usage reached"),
|
||||
CooldownRule(cooldown=timedelta(seconds=30), failure_message="Cooldown active"),
|
||||
TimeWindowRule(
|
||||
start_time=time(9, 0),
|
||||
end_time=time(17, 0),
|
||||
failure_message="Outside working hours"
|
||||
)
|
||||
])
|
||||
engine.set_command_rule("!test", rule)
|
||||
|
||||
# First usage during working hours should pass
|
||||
context = CommandContext(
|
||||
command="!test",
|
||||
user_id="user1",
|
||||
current_time=datetime(2024, 1, 15, 12, 0, 0)
|
||||
)
|
||||
should_trigger, message = engine.evaluate(context)
|
||||
assert should_trigger is True
|
||||
|
||||
# Record the usage
|
||||
engine.record_usage("!test", "user1", datetime(2024, 1, 15, 12, 0, 0))
|
||||
|
||||
# Same time should fail due to cooldown
|
||||
context = CommandContext(
|
||||
command="!test",
|
||||
user_id="user2",
|
||||
current_time=datetime(2024, 1, 15, 12, 0, 5)
|
||||
)
|
||||
should_trigger, message = engine.evaluate(context)
|
||||
assert should_trigger is False
|
||||
assert message == "Cooldown active"
|
||||
|
||||
# Outside working hours should fail
|
||||
context = CommandContext(
|
||||
command="!test",
|
||||
user_id="user1",
|
||||
current_time=datetime(2024, 1, 15, 20, 0, 0)
|
||||
)
|
||||
should_trigger, message = engine.evaluate(context)
|
||||
assert should_trigger is False
|
||||
assert message == "Outside working hours"
|
||||
|
||||
|
||||
class TestIntegration:
|
||||
"""Integration tests for realistic use cases."""
|
||||
|
||||
def test_rate_limited_command(self):
|
||||
"""Test a command with a 10-second rate limit."""
|
||||
engine = RuleEngine()
|
||||
engine.add_global_rule(
|
||||
GlobalCooldownRule(
|
||||
cooldown=timedelta(seconds=10),
|
||||
failure_message="Wait 10 seconds"
|
||||
)
|
||||
)
|
||||
|
||||
base_time = datetime(2024, 1, 15, 12, 0, 0)
|
||||
|
||||
# First command passes
|
||||
context = CommandContext(command="!cat", user_id="user1", current_time=base_time)
|
||||
should_trigger, _ = engine.evaluate(context)
|
||||
assert should_trigger is True
|
||||
engine.record_usage("!cat", "user1", base_time)
|
||||
|
||||
# 5 seconds later fails
|
||||
context = CommandContext(
|
||||
command="!dog",
|
||||
user_id="user2",
|
||||
current_time=base_time + timedelta(seconds=5)
|
||||
)
|
||||
should_trigger, message = engine.evaluate(context)
|
||||
assert should_trigger is False
|
||||
assert message == "Wait 10 seconds"
|
||||
|
||||
# 15 seconds later passes
|
||||
context = CommandContext(
|
||||
command="!dog",
|
||||
user_id="user2",
|
||||
current_time=base_time + timedelta(seconds=15)
|
||||
)
|
||||
should_trigger, _ = engine.evaluate(context)
|
||||
assert should_trigger is True
|
||||
|
||||
def test_reaction_gated_command(self):
|
||||
"""Test a command that requires 5 reactions to trigger."""
|
||||
engine = RuleEngine()
|
||||
engine.set_command_rule(
|
||||
"!special",
|
||||
MinReactionsRule(min_reactions=5, failure_message="Need 5 reactions")
|
||||
)
|
||||
|
||||
# 3 reactions - fails
|
||||
context = CommandContext(command="!special", user_id="user1", reaction_count=3)
|
||||
should_trigger, message = engine.evaluate(context)
|
||||
assert should_trigger is False
|
||||
assert message == "Need 5 reactions"
|
||||
|
||||
# 5 reactions - passes
|
||||
context = CommandContext(command="!special", user_id="user1", reaction_count=5)
|
||||
should_trigger, _ = engine.evaluate(context)
|
||||
assert should_trigger is True
|
||||
|
||||
def test_prevent_spam_from_same_user(self):
|
||||
"""Test preventing the same user from using a command twice in a row."""
|
||||
engine = RuleEngine()
|
||||
engine.set_command_rule(
|
||||
"!share",
|
||||
NotSameUserRule(failure_message="Wait for someone else")
|
||||
)
|
||||
|
||||
# First usage passes
|
||||
context = CommandContext(command="!share", user_id="user1")
|
||||
should_trigger, _ = engine.evaluate(context)
|
||||
assert should_trigger is True
|
||||
engine.record_usage("!share", "user1")
|
||||
|
||||
# Same user fails
|
||||
context = CommandContext(command="!share", user_id="user1")
|
||||
should_trigger, message = engine.evaluate(context)
|
||||
assert should_trigger is False
|
||||
assert message == "Wait for someone else"
|
||||
|
||||
# Different user passes
|
||||
context = CommandContext(command="!share", user_id="user2")
|
||||
should_trigger, _ = engine.evaluate(context)
|
||||
assert should_trigger is True
|
||||
|
||||
def test_limited_usage_command(self):
|
||||
"""Test a command with a maximum usage count."""
|
||||
engine = RuleEngine()
|
||||
engine.set_command_rule(
|
||||
"!prize",
|
||||
MaxUsageRule(max_usage=3, failure_message="Prize limit reached")
|
||||
)
|
||||
|
||||
# First 3 uses pass
|
||||
for i in range(3):
|
||||
context = CommandContext(command="!prize", user_id=f"user{i}")
|
||||
should_trigger, _ = engine.evaluate(context)
|
||||
assert should_trigger is True
|
||||
engine.record_usage("!prize", f"user{i}")
|
||||
|
||||
# 4th use fails
|
||||
context = CommandContext(command="!prize", user_id="user4")
|
||||
should_trigger, message = engine.evaluate(context)
|
||||
assert should_trigger is False
|
||||
assert message == "Prize limit reached"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
Loading…
Reference in New Issue
Block a user