MacOS Python script - Test external USB/SSD performance

import os
import time
import sys
import secrets

def get_size(bytes, suffix="B"):
    """Scale bytes to its proper format"""
    factor = 1024
    for unit in ["", "K", "M", "G", "T", "P"]:
        if bytes < factor:
            return f"{bytes:.2f}{unit}{suffix}"
        bytes /= factor

def run_test():
    print("--- macOS Python SSD Benchmark ---")
    path = input("Drag and drop the target folder here: ").strip().replace("'", "").replace('"', '').strip()

    if not os.path.isdir(path):
        print(f"Error: {path} is not a valid directory.")
        return

    filename = os.path.join(path, "py_test_file.bin")
    # File size: 1GB (adjust strictly if needed)
    file_size_mb = 1024
    file_size_bytes = file_size_mb * 1024 * 1024
    buffer_size = 1024 * 1024 # 1MB chunks

    print(f"\nTarget: {path}")
    print(f"File Size: {file_size_mb} MB")

    # Generate random data chunk (1MB) to defeat compression algorithms
    print("Preparing random data to ensure accuracy...")
    random_data = secrets.token_bytes(buffer_size)

    # --- WRITE TEST ---
    print(f"\n[1/3] Running Write Test...")
    try:
        start_time = time.time()
        with open(filename, 'wb') as f:
            for _ in range(file_size_mb):
                f.write(random_data)
            # Force write to disk (bypass OS write cache)
            f.flush()
            os.fsync(f.fileno())
        end_time = time.time()

        duration = end_time - start_time
        speed = (file_size_bytes / duration) / (1024 * 1024)
        print(f"✅ Write Speed: {speed:.2f} MB/s")

    except Exception as e:
        print(f"Write failed: {e}")
        return

    # --- FLUSH CACHE ---
    print("\n[2/3] Flushing RAM Cache...")
    print("NOTE: We need to purge macOS file cache to measure real disk Read speed.")
    print("You may be asked for your password (sudo purge).")
    os.system("sudo purge")
    # Wait a moment for system to settle
    time.sleep(2)

    # --- READ TEST ---
    print("\n[3/3] Running Read Test...")
    try:
        start_time = time.time()
        with open(filename, 'rb') as f:
            while True:
                data = f.read(buffer_size)
                if not data:
                    break
        end_time = time.time()

        duration = end_time - start_time
        speed = (file_size_bytes / duration) / (1024 * 1024)
        print(f"✅ Read Speed:  {speed:.2f} MB/s")

    except Exception as e:
        print(f"Read failed: {e}")

    # --- CLEANUP ---
    try:
        os.remove(filename)
        print("\nTest file removed.")
    except:
        print("\nCould not remove test file.")

if __name__ == "__main__":
    run_test()