Optimize imports for tests as well

Change-Id: I2124928c85cebe9d752bbca6847484b3d32d9fea