[Agalya/Preethi/DD] Fixing an edge case where entity is nil - Time to refactor