disk_benchmark.py

#!/usr/bin/env python3
"""
Disk Benchmark Tool
Let's test how fast your disk really is!
"""

import os
import sys
import time
import random
import shutil
from pathlib import Path

class DiskBenchmark:
    def __init__(self, test_dir=None, file_size_mb=100, num_files=5):
        """
        Set up everything we need for the benchmark
        
        Args:
            test_dir: Where should we run the tests? (defaults to where this script lives)
            file_size_mb: How big should each test file be? (in MB)
            num_files: How many files should we create for testing?
        """
        self.test_dir = Path(test_dir) if test_dir else Path(__file__).parent
        self.file_size_mb = file_size_mb
        self.num_files = num_files
        self.file_size_bytes = file_size_mb * 1024 * 1024
        self.test_data_dir = self.test_dir / "benchmark_temp"
        
    def format_speed(self, bytes_per_sec):
        """Turn raw bytes into something we can actually understand"""
        mb_per_sec = bytes_per_sec / (1024 * 1024)
        if mb_per_sec >= 1024:
            return f"{mb_per_sec / 1024:.2f} GB/s"
        else:
            return f"{mb_per_sec:.2f} MB/s"
    
    def format_size(self, bytes_size):
        """Make file sizes easier to read"""
        mb = bytes_size / (1024 * 1024)
        if mb >= 1024:
            return f"{mb / 1024:.2f} GB"
        else:
            return f"{mb:.2f} MB"
    
    def setup(self):
        """Get everything ready for testing"""
        print(f"📁 Where we're testing: {self.test_dir.absolute()}")
        print(f"📊 File size: {self.file_size_mb} MB")
        print(f"📈 Number of files: {self.num_files}")
        print(f"💾 Total data we'll write: {self.format_size(self.file_size_bytes * self.num_files)}")
        print("-" * 60)
        
        if self.test_data_dir.exists():
            shutil.rmtree(self.test_data_dir)
        self.test_data_dir.mkdir(exist_ok=True)
    
    def cleanup(self):
        """Clean up all the test files we created"""
        if self.test_data_dir.exists():
            shutil.rmtree(self.test_data_dir)
    
    def benchmark_sequential_write(self):
        """Let's see how fast we can write files in order"""
        print("\n🔹 Sequential Write Test")
        total_bytes = 0
        start_time = time.time()
        
        # Create some random data once so the CPU doesn't slow us down
        chunk_size = 1024 * 1024  # 1 MB chunks
        data = os.urandom(chunk_size)
        
        for i in range(self.num_files):
            file_path = self.test_data_dir / f"test_seq_{i}.dat"
            with open(file_path, 'wb') as f:
                bytes_written = 0
                while bytes_written < self.file_size_bytes:
                    f.write(data)
                    bytes_written += chunk_size
                    total_bytes += chunk_size
            print(f"  Written file {i+1}/{self.num_files}", end='\r')
        
        elapsed = time.time() - start_time
        speed = total_bytes / elapsed
        
        print(f"  ✓ All done in {elapsed:.2f} seconds")
        print(f"  Speed: {self.format_speed(speed)}")
        return speed
    
    def benchmark_sequential_read(self):
        """Now let's see how fast we can read those files back"""
        print("\n🔹 Sequential Read Test")
        total_bytes = 0
        start_time = time.time()
        
        chunk_size = 1024 * 1024  # 1 MB chunks
        
        for i in range(self.num_files):
            file_path = self.test_data_dir / f"test_seq_{i}.dat"
            with open(file_path, 'rb') as f:
                while True:
                    data = f.read(chunk_size)
                    if not data:
                        break
                    total_bytes += len(data)
            print(f"  Read file {i+1}/{self.num_files}", end='\r')
        
        elapsed = time.time() - start_time
        speed = total_bytes / elapsed
        
        print(f"  ✓ All done in {elapsed:.2f} seconds")
        print(f"  Speed: {self.format_speed(speed)}")
        return speed
    
    def benchmark_random_write(self):
        """Time to test writing data all over the place (like real-world usage)"""
        print("\n🔹 Random Write Test (4KB blocks)")
        
        file_path = self.test_data_dir / "test_random.dat"
        block_size = 4096  # 4 KB
        num_blocks = self.file_size_bytes // block_size
        
        # Set up the file first
        with open(file_path, 'wb') as f:
            f.seek(self.file_size_bytes - 1)
            f.write(b'\0')
        
        # Generate some random data
        data = os.urandom(block_size)
        
        start_time = time.time()
        with open(file_path, 'r+b') as f:
            for i in range(num_blocks):
                offset = random.randint(0, num_blocks - 1) * block_size
                f.seek(offset)
                f.write(data)
                if i % 1000 == 0:
                    print(f"  Written {i}/{num_blocks} blocks", end='\r')
        
        elapsed = time.time() - start_time
        speed = self.file_size_bytes / elapsed
        
        print(f"  ✓ All done in {elapsed:.2f} seconds")
        print(f"  Speed: {self.format_speed(speed)}")
        return speed
    
    def benchmark_random_read(self):
        """And now let's read data randomly (like your computer does all day)"""
        print("\n🔹 Random Read Test (4KB blocks)")
        
        file_path = self.test_data_dir / "test_random.dat"
        block_size = 4096  # 4 KB
        num_blocks = self.file_size_bytes // block_size
        
        start_time = time.time()
        with open(file_path, 'rb') as f:
            for i in range(num_blocks):
                offset = random.randint(0, num_blocks - 1) * block_size
                f.seek(offset)
                data = f.read(block_size)
                if i % 1000 == 0:
                    print(f"  Read {i}/{num_blocks} blocks", end='\r')
        
        elapsed = time.time() - start_time
        speed = self.file_size_bytes / elapsed
        
        print(f"  ✓ All done in {elapsed:.2f} seconds")
        print(f"  Speed: {self.format_speed(speed)}")
        return speed
    
    def run(self):
        """Let's run all the tests and see what your disk can do!"""
        print("=" * 60)
        print("🚀 DISK BENCHMARK TOOL")
        print("=" * 60)
        
        try:
            self.setup()
            
            results = {}
            results['seq_write'] = self.benchmark_sequential_write()
            results['seq_read'] = self.benchmark_sequential_read()
            results['rand_write'] = self.benchmark_random_write()
            results['rand_read'] = self.benchmark_random_read()
            
            print("\n" + "=" * 60)
            print("📊 HERE'S WHAT WE FOUND")
            print("=" * 60)
            print(f"Sequential Write: {self.format_speed(results['seq_write'])}")
            print(f"Sequential Read:  {self.format_speed(results['seq_read'])}")
            print(f"Random Write:     {self.format_speed(results['rand_write'])}")
            print(f"Random Read:      {self.format_speed(results['rand_read'])}")
            print("=" * 60)
            
        except KeyboardInterrupt:
            print("\n\n⚠️  Alright, stopping the benchmark...")
        except Exception as e:
            print(f"\n\n❌ Oops, something went wrong: {e}")
        finally:
            print("\n🧹 Cleaning up the mess we made...")
            self.cleanup()
            print("✓ All done!")

def main():
    import argparse
    
    parser = argparse.ArgumentParser(
        description='Disk Benchmark Tool - Let\'s see how fast your disk really is!',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s                          # Run with defaults (100MB, 5 files)
  %(prog)s --size 500 --files 10    # Test with 500MB files, 10 files
  %(prog)s --dir /path/to/test      # Test in specific directory
        """
    )
    
    parser.add_argument(
        '--dir',
        type=str,
        help='Directory to run tests in (default: current directory)',
        default=None
    )
    
    parser.add_argument(
        '--size',
        type=int,
        help='Size of each test file in MB (default: 100)',
        default=100
    )
    
    parser.add_argument(
        '--files',
        type=int,
        help='Number of files to test with (default: 5)',
        default=5
    )
    
    args = parser.parse_args()
    
    benchmark = DiskBenchmark(
        test_dir=args.dir,
        file_size_mb=args.size,
        num_files=args.files
    )
    
    benchmark.run()

if __name__ == '__main__':
    main()