Add tests and fix errors for generator

Added integration tests to verify that the generator was broken and not to break again, then fixes #10 and #11.

I can rebase #6 after this pull request is merged, if you like.