Drime upload and download testing script

drm-test.py

#!/usr/bin/env python3

import logging
import os
import subprocess as sp
import sys
import time
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from hashlib import md5


TESTS = {
    '5K': 5 * 1024,
    '5M': 5 * 1024**2,
    '508M': 508 * 1024**2,
    '1.05G': int(1.05 * 1024**3),
    '5.2G': int(5.2 * 1024**3),
    #'10.1G': int(10.1 * 1024**3),
    #'50G': 50 * 1024**3,
}
BUFSIZE = 4 * 1024  # 4k

def get_args():
    rclone = os.path.join(os.path.dirname(__file__), 'rclone')
    remote_opts = ('drm_secret', 'drime_test')
    p = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
    p.add_argument('-r', '--rclone', default=rclone, help='The path to the '
        'rclone binary to use')
    p.add_argument('-m', '--remote', default='drm_secret', choices=remote_opts,
        help='Specify which remote to use.  The default is the encrypted '
        'variant of the remote')
    p.add_argument('-c', '--cleanup', action='store_true', default=False,
        help='Remove all the files after the testing has concluded')
    p.add_argument('-D', '--debug', action='store_true', default=False,
        help='Add debug output')

    args = p.parse_args()

    return args


def setup_logging(args):
    level = logging.DEBUG if args.debug else logging.INFO
    logging.basicConfig(
        format=(
            '%(asctime)s - %(levelname)s - '
            '%(filename)s:%(lineno)d %(funcName)s - %(message)s'
        ),
        level=level,
    )


def get_md5(bin_fname):
    md5sum = md5()
    with open(bin_fname, 'rb') as fh:
        while True:
            buf = fh.read(BUFSIZE)
            md5sum.update(buf)

            if len(buf) != BUFSIZE:
                break

    return md5sum.hexdigest()


def get_md5_from_f(md5name):
    with open(md5name) as fh:
        return fh.read().strip()


def create_file(fname, size):
    with open('/dev/urandom', 'rb') as urfh:
        with open(fname, 'wb') as bfh:
            md5sum = md5()
            bytes_written = 0

            while True:
                # Read either the full bufsize or the remaining bytes
                to_read = min(BUFSIZE, size - bytes_written)
                buf = urfh.read(to_read)
                # Calculate the md5 sum while we're reading/writing
                md5sum.update(buf)
                bfh.write(buf)

                bytes_written += to_read

                # Have we written the entire target size?
                if bytes_written >= size:
                    break

            # Write out the md5 to .md5 file
            with open(f'{fname}.md5', 'w') as fh:
                fh.write(md5sum.hexdigest())

            return md5sum.hexdigest()


def write_md5hex(md5sum, fname):
    with open(fname, 'w') as fh:
        fh.write(md5sum)


def upload_file(fname, rclone, remote):
    cmd = [rclone, 'copy', fname, f'{remote}:/']

    logging.info(f'Running: {" ".join(cmd)}')
    p = sp.run(cmd)

    if p.returncode != 0:
        logging.error(f'Rclone exited with {p.returncode}')


def download_file(fname, rclone, remote):
    dl = f'{fname}.dl'
    cmd = [rclone, 'copyto', f'{remote}:/{fname}', dl]

    logging.info(f'Running: {" ".join(cmd)}')
    p = sp.run(cmd)

    if p.returncode != 0:
        logging.error(f'Rclone exited with {p.returncode}')

    return dl


def get_mbps(size_bytes, duration):
    # Calculate bits/sec
    bps = size_bytes * 8 / duration
    # Turn it into Mb/s
    mbps = bps / 1024**2

    return mbps


def run_test(tsize, size, args):
    bfname = f'{tsize}.bin'
    md5name = f'{bfname}.md5'
    if os.path.isfile(bfname) and os.path.getsize(bfname) == size:
        md5hex = get_md5(bfname)
        try:
            md5sum = get_md5_from_f(md5name)
        except Exception as e:
            logging.error(f'Failed to get md5 from {md5name}: {e}')
            # There's a problem reading the md5 file, let's write it out
            write_md5hex(md5hex, md5name)
        else:
            if md5hex != md5sum:
                logging.warning('MD5 sums do not match, writing newly '
                                f'calculated sum to {md5name}')
                write_md5hex(md5hex, md5name)

        logging.warning(f'{bfname} exists and is the correct size/md5, '
                        'skipping creation')
    else:
        md5hex = create_file(bfname, size)

    start = time.time()
    upload_file(bfname, args.rclone, args.remote)
    duration = time.time() - start
    mbps = get_mbps(size, duration)

    logging.info(f'Upload completed in {duration:.02f} seconds at '
                 f'{mbps:.02f} Mb/s')

    start = time.time()
    dl_name = download_file(bfname, args.rclone, args.remote)
    duration = time.time() - start
    mbps = get_mbps(size, duration)
    logging.info(f'Download completed in {duration:.02f} seconds at '
                 f'{mbps:.02f} Mb/s')

    dl_md5 = get_md5(dl_name)

    if dl_md5 == md5hex:
        logging.info(f'Upload and download of {tsize} file matches!')
    else:
        logging.error(f'Upload and download of {tsize} file does not '
                      'match the md5 sums')

    logging.info(f'Cleaning up the downloaded file: {dl_name}')
    os.unlink(dl_name)

    if args.cleanup:
        logging.info(f'Removing {bfname} and {md5name}')
        os.unlink(bfname)
        os.unlink(md5name)


def main():
    args = get_args()
    setup_logging(args)

    for tsize, size in TESTS.items():
        run_test(tsize, size, args)

    return 0


if __name__ == '__main__':
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        sys.exit(0)