Fix to Unit Tests