Add load_gen.py tests

@karlht can you review