Updated test_generator unit test to match new structure.