Fix failing tests/code