Small fixes
Fix failing tests