Fixing naive bayes unit tests