[#6] Test fixes