Fixed some flaky tests