95 lines
3.5 KiB
Python
95 lines
3.5 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
|
|
|
|
log_file = os.path.join(temp_dir, "big.log")
|
|
with open(log_file, "w") as f:
|
|
f.write("A" * 1000)
|
|
|
|
print(f"Testing warning threshold. State file: {state_file}")
|
|
|
|
# Mock send_email internals
|
|
with patch('smtplib.SMTP') as mock_smtp:
|
|
|
|
# --- Scenario 1: Warning Level (85%) ---
|
|
print("\n--- Scenario 1: Warning Level (85%) ---")
|
|
|
|
# Mock Usage: 85%
|
|
disk_cleaner.shutil.disk_usage = MagicMock(return_value=MagicMock(total=1000, used=850, free=150))
|
|
|
|
# Mock Partitions
|
|
disk_cleaner.os.path.exists = lambda p: True
|
|
disk_cleaner.get_partitions = MagicMock(return_value=[('/dev/test', temp_dir)])
|
|
|
|
# Mock Open Files
|
|
disk_cleaner.get_open_files_flat = MagicMock(return_value=[
|
|
{'path': log_file, 'pid': 1234, 'size': 1024 * 1024 * 50} # 50MB
|
|
])
|
|
|
|
# Run
|
|
disk_cleaner.check_and_clean()
|
|
|
|
# Verify Email: Context Manager Mock Pattern
|
|
smtp_instance = mock_smtp.return_value.__enter__.return_value
|
|
|
|
if smtp_instance.send_message.called:
|
|
args, _ = smtp_instance.send_message.call_args
|
|
msg = args[0]
|
|
if "WARNING:" in msg['Subject']:
|
|
print("SUCCESS: Warning email sent.")
|
|
print(f"Subject: {msg['Subject']}")
|
|
if "Suspected largest open file" in msg.get_content():
|
|
print("SUCCESS: Body identifies suspect file.")
|
|
else:
|
|
print("FAILURE: Body missing suspect file.")
|
|
else:
|
|
print(f"FAILURE: Wrong subject: {msg['Subject']}")
|
|
else:
|
|
print("FAILURE: No email sent.")
|
|
|
|
mock_smtp.reset_mock()
|
|
|
|
# Reset state for next test
|
|
os.remove(state_file)
|
|
|
|
# --- Scenario 2: Critical Level (96%) ---
|
|
print("\n--- Scenario 2: Critical Level (96%) ---")
|
|
|
|
# Mock Usage: 96%
|
|
disk_cleaner.shutil.disk_usage = MagicMock(return_value=MagicMock(total=1000, used=960, free=40))
|
|
|
|
# Run
|
|
disk_cleaner.check_and_clean()
|
|
|
|
# Verify Email
|
|
smtp_instance = mock_smtp.return_value.__enter__.return_value
|
|
|
|
if smtp_instance.send_message.called:
|
|
args, _ = smtp_instance.send_message.call_args
|
|
msg = args[0]
|
|
if "URGENT:" in msg['Subject']:
|
|
print("SUCCESS: Urgent email sent.")
|
|
else:
|
|
print(f"FAILURE: Wrong subject: {msg['Subject']}")
|
|
else:
|
|
pass
|
|
|
|
if __name__ == "__main__":
|
|
# Mock is_log_file to be true so Critical clean attempts it
|
|
disk_cleaner.is_log_file = MagicMock(return_value=(True, "Yes"))
|
|
|
|
# Mock shrink_file_inplace to succeed
|
|
disk_cleaner.shrink_file_inplace = MagicMock(return_value=(True, "Shrunk"))
|
|
|
|
run_test()
|