Update tests, fix some errors in main script