#!/usr/bin/env python3

import sys

import zst_zstd
import zst_zstandard
import zst_pyzstd


def small_tests(module):
    """main function"""
    all_good = True
    with open('/etc/protocols', 'rb') as protocols_file:
        protocols = protocols_file.read()
    for string in [b'Rush Rocks', b'flurfl', protocols]:
        compressed = module.compress(string)
        decompressed = module.decompress(compressed)
        if string != decompressed:
            all_good = False
            sys.stderr.write('%s: Bad round trip: "%s", "%s"\n' % (sys.argv[0], string[:40], decompressed[:40]))
    return all_good


def main():
    """Run the tests, once for each module."""
    all_good = True
    for module in [zst_zstd, zst_zstandard, zst_pyzstd]:
        all_good &= small_tests(module)
    if all_good:
        print('All tests passed')
    else:
        raise SystemExit('One or more tests failed')


if __name__ == '__main__':
    main()