Fix forgotten imports after refactoring

Change-Id: I546875c471dbf49fe556da3d4a884bff6792706b
Merge "Fix forgotten imports after refactoring"