63 lines
2.2 KiB
Python
63 lines
2.2 KiB
Python
import unittest
|
|
from unittest.mock import MagicMock, patch, mock_open
|
|
import os
|
|
import time
|
|
import json
|
|
import disk_cleaner
|
|
import tempfile
|
|
import sys
|
|
|
|
def run_test():
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
state_file = os.path.join(temp_dir, "cleaner_state.json")
|
|
disk_cleaner.STATE_FILE = state_file
|
|
|
|
print(f"Testing rate limiting. State file: {state_file}")
|
|
|
|
# Mock send_email internals (smtplib) to verify call count
|
|
with patch('smtplib.SMTP') as mock_smtp:
|
|
# 1. First Call: Should Send
|
|
print("First Trigger (expect send)...")
|
|
disk_cleaner.send_email("Subject 1", "Body", "/mnt/data")
|
|
|
|
if mock_smtp.called:
|
|
print("SUCCESS: Email sent.")
|
|
else:
|
|
print("FAILURE: Email NOT sent.")
|
|
mock_smtp.reset_mock()
|
|
|
|
# Verify state file
|
|
if os.path.exists(state_file):
|
|
with open(state_file) as f:
|
|
state = json.load(f)
|
|
print(f"State: {state}")
|
|
if '/mnt/data' in state:
|
|
print("SUCCESS: State recorded.")
|
|
else:
|
|
print("FAILURE: State missing key.")
|
|
else:
|
|
print("FAILURE: State file not created.")
|
|
|
|
# 2. Second Call (Immediate): Should Suppress
|
|
print("Second Trigger (Immediate - expect suppress)...")
|
|
disk_cleaner.send_email("Subject 2", "Body", "/mnt/data")
|
|
|
|
if not mock_smtp.called:
|
|
print("SUCCESS: Email suppressed.")
|
|
else:
|
|
print("FAILURE: Email sent despite cooldown.")
|
|
mock_smtp.reset_mock()
|
|
|
|
# 3. Third Call (Different Mount): Should Send
|
|
print("Third Trigger (Different Mount - expect send)...")
|
|
disk_cleaner.send_email("Subject 3", "Body", "/mnt/other")
|
|
|
|
if mock_smtp.called:
|
|
print("SUCCESS: Email sent (different key).")
|
|
else:
|
|
print("FAILURE: Email NOT sent for new key.")
|
|
mock_smtp.reset_mock()
|
|
|
|
if __name__ == "__main__":
|
|
run_test()
|