Fix integration tests

The integration tests are failing due to the change to astropy>=5.

This MR fixes the issue and adds a new test that the generation function works on DataFrame inputs.