Fixing test error. Adding import.

(cherry picked from commit b68cdd1ffd98ca226d5647916c117b68c9de243b)