Cleanup tests to use unittest framework