Refactor unit test cases for EntityTests

Added some assertions and refactored a couple test cases