diff --git a/api/test/features_servicer_test.py b/api/test/features_servicer_test.py
index 177bc3f..21154de 100644
--- a/api/test/features_servicer_test.py
+++ b/api/test/features_servicer_test.py
@@ -9,7 +9,10 @@
 from __future__ import absolute_import
 
 import unittest
-import mox
+try:
+  from mox3 import mox
+except ImportError:
+  import mox
 
 from components.prpc import codes
 from components.prpc import context