diff --git a/api/controller/image_unittest.py b/api/controller/image_unittest.py
index d47a3bb..824b9bf 100644
--- a/api/controller/image_unittest.py
+++ b/api/controller/image_unittest.py
@@ -10,7 +10,7 @@
 import mock
 import os
 
-from chromite.api.gen import image_pb2
+from chromite.api.gen.chromite.api import image_pb2
 from chromite.api.controller import image as image_controller
 from chromite.lib import constants
 from chromite.lib import cros_test_lib