Improve test robustness