diff --git a/test/unit/extensions/omci/test_image_agent.py b/test/unit/extensions/omci/test_image_agent.py
index 08f0f1e..0e500cb 100644
--- a/test/unit/extensions/omci/test_image_agent.py
+++ b/test/unit/extensions/omci/test_image_agent.py
@@ -26,8 +26,7 @@
         OmciDownloadSection, OmciDownloadSectionLast, OmciDownloadSectionResponse, \
         OmciActivateImage, OmciActivateImageResponse,  \
         OmciCommitImage, OmciCommitImageResponse
-from pyvoltha.protos.voltha_pb2 import ImageDownload
-from pyvoltha.protos.device_pb2 import Device
+from voltha_protos.device_pb2 import Device, ImageDownload
 
 from test.unit.extensions.omci.mock.mock_adapter_agent import MockAdapterAgent, MockCore
 from twisted.internet import reactor