Fix broken imports in tests

The imports in some tests seem to have been broken in commit 1bf247f6ecfceb4722bf7a11d953ca8d70252b98.