Changed entity class and updated unit test
Completed merge, tests pass YAY