Fixed test data generation

Now it generates the right test set