Fix NPEs and bugs in unit testing