Refactor generators and add more tests