diff --git a/au_test_harness/gce.py b/au_test_harness/gce.py
new file mode 100644
index 0000000..2427878
--- /dev/null
+++ b/au_test_harness/gce.py
@@ -0,0 +1,716 @@
+# -*- coding: utf-8 -*-
+# Copyright 2015 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""A convinient wrapper of the GCE python API.
+
+Public methods in class GceContext raise HttpError when the underlining call to
+Google API fails, or gce.Error on other failures.
+"""
+
+from __future__ import print_function
+
+import sys
+
+from googleapiclient.discovery import build
+from googleapiclient.errors import HttpError
+from googleapiclient.http import HttpRequest
+import httplib2
+from oauth2client.client import GoogleCredentials
+
+from chromite.lib import cros_logging as logging
+from chromite.lib import timeout_util
+
+
+class Error(Exception):
+  """Base exception for this module."""
+
+
+class ResourceNotFoundError(Error):
+  """Exceptions raised when requested GCE resource was not found."""
+
+
+class RetryOnServerErrorHttpRequest(HttpRequest):
+  """A HttpRequest that will be retried on server errors automatically."""
+
+  def __init__(self, num_retries, *args, **kwargs):
+    """Constructor for RetryOnServerErrorHttpRequest."""
+    self.num_retries = num_retries
+    super(RetryOnServerErrorHttpRequest, self).__init__(*args, **kwargs)
+
+  def execute(self, http=None, num_retries=None):
+    """Excutes a RetryOnServerErrorHttpRequest.
+
+    HttpRequest.execute() has the option of automatically retrying on server
+    errors, i.e., 500 status codes. Call it with a non-zero value of
+    |num_retries| will cause failed requests to be retried.
+
+    Args:
+      http: The httplib2.http to send this request through.
+      num_retries: Number of retries. Class default value will be used if
+          omitted.
+
+    Returns:
+      A deserialized object model of the response body as determined
+          by the postproc. See HttpRequest.execute().
+    """
+    return super(RetryOnServerErrorHttpRequest, self).execute(
+        http=http, num_retries=num_retries or self.num_retries)
+
+
+def _GetMetdataValue(metadata, key):
+  """Finds a value corresponding to a given metadata key.
+
+  Args:
+    metadata: metadata object, i.e. a dict containing containing 'items'
+      - a list of key-value pairs.
+    key: name of the key.
+
+  Returns:
+    Corresponding value or None if it was not found.
+  """
+  for item in metadata['items']:
+    if item['key'] == key:
+      return item['value']
+  return None
+
+
+def _UpdateMetadataValue(metadata, key, value):
+  """Updates a single key-value pair in a metadata object.
+
+  Args:
+    metadata: metadata object, i.e. a dict containing containing 'items'
+      - a list of key-value pairs.
+    key: name of the key.
+    value: new value for the key, or None if it should be removed.
+  """
+  items = metadata.setdefault('items', [])
+  for item in items:
+    if item['key'] == key:
+      if value is None:
+        items.remove(item)
+      else:
+        item['value'] = value
+      return
+
+  if value is not None:
+    items.append({
+        'key': key,
+        'value': value,
+    })
+
+
+class GceContext(object):
+  """A convinient wrapper around the GCE Python API."""
+
+  # These constants are made public so that users can customize as they need.
+  DEFAULT_TIMEOUT_SEC = 5 * 60
+  INSTANCE_OPERATIONS_TIMEOUT_SEC = 10 * 60
+  IMAGE_OPERATIONS_TIMEOUT_SEC = 10 * 60
+
+  _GCE_SCOPES = (
+      'https://www.googleapis.com/auth/compute',  # CreateInstance, CreateImage
+      'https://www.googleapis.com/auth/devstorage.full_control', # CreateImage
+  )
+  _DEFAULT_NETWORK = 'default'
+  _DEFAULT_MACHINE_TYPE = 'n1-standard-8'
+
+  # Project default service account and scopes.
+  _DEFAULT_SERVICE_ACCOUNT_EMAIL = 'default'
+  # The list is in line with what the gcloud cli uses.
+  # https://cloud.google.com/sdk/gcloud/reference/compute/instances/create
+  _DEFAULT_INSTANCE_SCOPES = [
+      'https://www.googleapis.com/auth/cloud.useraccounts.readonly',
+      'https://www.googleapis.com/auth/devstorage.read_only',
+      'https://www.googleapis.com/auth/logging.write',
+  ]
+
+  # This is made public to allow easy customization of the retry behavior.
+  RETRIES = 2
+
+  def __init__(self, project, zone, credentials, thread_safe=False):
+    """Initializes GceContext.
+
+    Args:
+      project: The GCP project to create instances in.
+      zone: The default zone to create instances in.
+      credentials: The credentials used to call the GCE API.
+      thread_safe: Whether the client is expected to be thread safe.
+    """
+    self.project = project
+    self.zone = zone
+
+    def _BuildRequest(http, *args, **kwargs):
+      """Custom request builder."""
+      return self._BuildRetriableRequest(self.RETRIES, http, thread_safe,
+                                         credentials, *args, **kwargs)
+
+    self.gce_client = build('compute', 'v1', credentials=credentials,
+                            requestBuilder=_BuildRequest)
+
+    self.region = self.GetZoneRegion(zone)
+
+  @classmethod
+  def ForServiceAccount(cls, project, zone, json_key_file):
+    """Creates a GceContext using service account credentials.
+
+    About service account:
+    https://developers.google.com/api-client-library/python/auth/service-accounts
+
+    Args:
+      project: The GCP project to create images and instances in.
+      zone: The default zone to create instances in.
+      json_key_file: Path to the service account JSON key.
+
+    Returns:
+      GceContext.
+    """
+    credentials = GoogleCredentials.from_stream(json_key_file).create_scoped(
+        cls._GCE_SCOPES)
+    return GceContext(project, zone, credentials)
+
+  @classmethod
+  def ForServiceAccountThreadSafe(cls, project, zone, json_key_file):
+    """Creates a thread-safe GceContext using service account credentials.
+
+    About service account:
+    https://developers.google.com/api-client-library/python/auth/service-accounts
+
+    Args:
+      project: The GCP project to create images and instances in.
+      zone: The default zone to create instances in.
+      json_key_file: Path to the service account JSON key.
+
+    Returns:
+      GceContext.
+    """
+    credentials = GoogleCredentials.from_stream(json_key_file).create_scoped(
+        cls._GCE_SCOPES)
+    return GceContext(project, zone, credentials, thread_safe=True)
+
+  def CreateAddress(self, name, region=None):
+    """Reserves an external IP address.
+
+    Args:
+      name: The name to assign to the address.
+      region: Region to reserved the address in.
+
+    Returns:
+      The reserved address as a string.
+    """
+    body = {
+        'name': name,
+    }
+    operation = self.gce_client.addresses().insert(
+        project=self.project,
+        region=region or self.region,
+        body=body).execute()
+    self._WaitForRegionOperation(
+        operation['name'], region,
+        timeout_sec=self.INSTANCE_OPERATIONS_TIMEOUT_SEC)
+
+    address = self.gce_client.addresses().get(
+        project=self.project,
+        region=region or self.region,
+        address=name).execute()
+
+    return address['address']
+
+  def DeleteAddress(self, name, region=None):
+    """Frees up an external IP address.
+
+    Args:
+      name: The name of the address.
+      region: Region of the address.
+    """
+    operation = self.gce_client.addresses().delete(
+        project=self.project,
+        region=region or self.region,
+        address=name).execute()
+    self._WaitForRegionOperation(
+        operation['name'], region=region or self.region,
+        timeout_sec=self.INSTANCE_OPERATIONS_TIMEOUT_SEC)
+
+  def GetZoneRegion(self, zone=None):
+    """Resolves name of the region that a zone belongs to.
+
+    Args:
+      zone: The zone to resolve.
+
+    Returns:
+      Name of the region corresponding to the zone.
+    """
+    zone_resource = self.gce_client.zones().get(
+        project=self.project,
+        zone=zone or self.zone).execute()
+    return zone_resource['region'].split('/')[-1]
+
+  def CreateInstance(self, name, image, zone=None, network=None, subnet=None,
+                     machine_type=None, default_scopes=True,
+                     static_address=None, **kwargs):
+    """Creates an instance with the given image and waits until it's ready.
+
+    Args:
+      name: Instance name.
+      image: Fully spelled URL of the image, e.g., for private images,
+          'global/images/my-private-image', or for images from a
+          publicly-available project,
+          'projects/debian-cloud/global/images/debian-7-wheezy-vYYYYMMDD'.
+          Details:
+          https://cloud.google.com/compute/docs/reference/latest/instances/insert
+      zone: The zone to create the instance in. Default zone will be used if
+          omitted.
+      network: An existing network to create the instance in. Default network
+          will be used if omitted.
+      subnet: The subnet to create the instance in.
+      machine_type: The machine type to use. Default machine type will be used
+          if omitted.
+      default_scopes: If true, the default scopes are added to the instances.
+      static_address: External IP address to assign to the instance as a string.
+          If None an emphemeral address will be used.
+      kwargs: Other possible Instance Resource properties.
+          https://cloud.google.com/compute/docs/reference/latest/instances#resource
+          Note that values from kwargs will overrule properties constructed from
+          positinal arguments, i.e., name, image, zone, network and
+          machine_type.
+
+    Returns:
+      URL to the created instance.
+    """
+    logging.info('Creating instance "%s" with image "%s" ...', name, image)
+    network = 'global/networks/%s' % network or self._DEFAULT_NETWORK
+    machine_type = 'zones/%s/machineTypes/%s' % (
+        zone or self.zone, machine_type or self._DEFAULT_MACHINE_TYPE)
+    service_accounts = (
+        {
+            'email': self._DEFAULT_SERVICE_ACCOUNT_EMAIL,
+            'scopes': self._DEFAULT_INSTANCE_SCOPES,
+        },
+    ) if default_scopes else ()
+
+    config = {
+        'name': name,
+        'machineType': machine_type,
+        'disks': (
+            {
+                'boot': True,
+                'autoDelete': True,
+                'initializeParams': {
+                    'sourceImage': image,
+                },
+            },
+        ),
+        'networkInterfaces': (
+            {
+                'network': network,
+                'accessConfigs': (
+                    {
+                        'type': 'ONE_TO_ONE_NAT',
+                        'name': 'External NAT',
+                    },
+                ),
+            },
+        ),
+        'serviceAccounts' : service_accounts,
+    }
+    config.update(**kwargs)
+    if static_address is not None:
+      config['networkInterfaces'][0]['accessConfigs'][0]['natIP'] = (
+          static_address)
+    if subnet is not None:
+      region = self.GetZoneRegion(zone)
+      config['networkInterfaces'][0]['subnetwork'] = (
+          'regions/%s/subnetworks/%s' % (region, subnet)
+      )
+    operation = self.gce_client.instances().insert(
+        project=self.project,
+        zone=zone or self.zone,
+        body=config).execute()
+    self._WaitForZoneOperation(
+        operation['name'],
+        timeout_sec=self.INSTANCE_OPERATIONS_TIMEOUT_SEC,
+        timeout_handler=lambda: self.DeleteInstance(name))
+    return operation['targetLink']
+
+  def DeleteInstance(self, name, zone=None):
+    """Deletes an instance with the name and waits until it's done.
+
+    Args:
+      name: Name of the instance to delete.
+      zone: Zone where the instance is in. Default zone will be used if omitted.
+    """
+    logging.info('Deleting instance "%s" ...', name)
+    operation = self.gce_client.instances().delete(
+        project=self.project,
+        zone=zone or self.zone,
+        instance=name).execute()
+    self._WaitForZoneOperation(
+        operation['name'], timeout_sec=self.INSTANCE_OPERATIONS_TIMEOUT_SEC)
+
+  def StartInstance(self, name, zone=None):
+    """Starts an instance with the name and waits until it's done.
+
+    Args:
+      name: Name of the instance to start.
+      zone: Zone where the instance is in. Default zone will be used if omitted.
+    """
+    logging.info('Starting instance "%s" ...', name)
+    operation = self.gce_client.instances().start(
+        project=self.project,
+        zone=zone or self.zone,
+        instance=name).execute()
+    self._WaitForZoneOperation(
+        operation['name'], timeout_sec=self.INSTANCE_OPERATIONS_TIMEOUT_SEC)
+
+  def StopInstance(self, name, zone=None):
+    """Stops an instance with the name and waits until it's done.
+
+    Args:
+      name: Name of the instance to stop.
+      zone: Zone where the instance is in. Default zone will be used if omitted.
+    """
+    logging.info('Stopping instance "%s" ...', name)
+    operation = self.gce_client.instances().stop(
+        project=self.project,
+        zone=zone or self.zone,
+        instance=name).execute()
+    self._WaitForZoneOperation(
+        operation['name'], timeout_sec=self.INSTANCE_OPERATIONS_TIMEOUT_SEC)
+
+  def CreateImage(self, name, source):
+    """Creates an image with the given |source|.
+
+    Args:
+      name: Name of the image to be created.
+      source:
+        Google Cloud Storage object of the source disk, e.g.,
+        'https://storage.googleapis.com/my-gcs-bucket/test_image.tar.gz'.
+
+    Returns:
+      URL to the created image.
+    """
+    logging.info('Creating image "%s" with "source" %s ...', name, source)
+    config = {
+        'name': name,
+        'rawDisk': {
+            'source': source,
+        },
+    }
+    operation = self.gce_client.images().insert(
+        project=self.project,
+        body=config).execute()
+    self._WaitForGlobalOperation(operation['name'],
+                                 timeout_sec=self.IMAGE_OPERATIONS_TIMEOUT_SEC,
+                                 timeout_handler=lambda: self.DeleteImage(name))
+    return operation['targetLink']
+
+  def DeleteImage(self, name):
+    """Deletes an image and waits until it's deleted.
+
+    Args:
+      name: Name of the image to delete.
+    """
+    logging.info('Deleting image "%s" ...', name)
+    operation = self.gce_client.images().delete(
+        project=self.project,
+        image=name).execute()
+    self._WaitForGlobalOperation(operation['name'],
+                                 timeout_sec=self.IMAGE_OPERATIONS_TIMEOUT_SEC)
+
+  def ListInstances(self, zone=None):
+    """Lists all instances.
+
+    Args:
+      zone: Zone where the instances are in. Default zone will be used if
+            omitted.
+
+    Returns:
+      A list of Instance Resources if found, or an empty list otherwise.
+    """
+    result = self.gce_client.instances().list(project=self.project,
+                                              zone=zone or self.zone).execute()
+    return result.get('items', [])
+
+  def ListImages(self):
+    """Lists all images.
+
+    Returns:
+      A list of Image Resources if found, or an empty list otherwise.
+    """
+    result = self.gce_client.images().list(project=self.project).execute()
+    return result.get('items', [])
+
+  def GetInstance(self, instance, zone=None):
+    """Gets an Instance Resource by name and zone.
+
+    Args:
+      instance: Name of the instance.
+      zone: Zone where the instance is in. Default zone will be used if omitted.
+
+    Returns:
+      An Instance Resource.
+
+    Raises:
+      ResourceNotFoundError if instance was not found, or HttpError on other
+      HTTP failures.
+    """
+    try:
+      return self.gce_client.instances().get(project=self.project,
+                                             zone=zone or self.zone,
+                                             instance=instance).execute()
+    except HttpError as e:
+      if e.resp.status == 404:
+        raise ResourceNotFoundError(
+            'Instance "%s" for project "%s" in zone "%s" was not found.' %
+            (instance, self.project, zone or self.zone))
+      else:
+        raise
+
+  def GetInstanceIP(self, instance, zone=None):
+    """Gets the external IP of an instance.
+
+    Args:
+      instance: Name of the instance to get IP for.
+      zone: Zone where the instance is in. Default zone will be used if omitted.
+
+    Returns:
+      External IP address of the instance.
+
+    Raises:
+      Error: Something went wrong when trying to get IP for the instance.
+    """
+    result = self.GetInstance(instance, zone)
+    try:
+      return result['networkInterfaces'][0]['accessConfigs'][0]['natIP']
+    except (KeyError, IndexError):
+      raise Error('Failed to get IP address for instance %s' % instance)
+
+  def GetInstanceInternalIP(self, instance, zone=None):
+    """Gets the internal IP of an instance."""
+    result = self.GetInstance(instance, zone)
+    try:
+      return result['networkInterfaces'][0]['networkIP']
+    except (KeyError, IndexError):
+      raise Error('Failed to get internal IP for instance %s' % instance)
+
+  def GetImage(self, image):
+    """Gets an Image Resource by name.
+
+    Args:
+      image: Name of the image to look for.
+
+    Returns:
+      An Image Resource.
+
+    Raises:
+      ResourceNotFoundError: The requested image was not found.
+    """
+    try:
+      return self.gce_client.images().get(project=self.project,
+                                          image=image).execute()
+    except HttpError as e:
+      if e.resp.status == 404:
+        raise ResourceNotFoundError('Image "%s" for project "%s" was not found.'
+                                    % (image, self.project))
+      else:
+        raise
+
+  def InstanceExists(self, instance, zone=None):
+    """Checks if an instance exists in the current project.
+
+    Args:
+      instance: Name of the instance to check existence of.
+      zone: Zone where the instance is in. Default zone will be used if omitted.
+
+    Returns:
+      True if the instance exists or False otherwise.
+    """
+    try:
+      return self.GetInstance(instance, zone) is not None
+    except ResourceNotFoundError:
+      return False
+
+  def ImageExists(self, image):
+    """Checks if an image exists in the current project.
+
+    Args:
+      image: Name of the image to check existence of.
+
+    Returns:
+      True if the instance exists or False otherwise.
+    """
+    try:
+      return self.GetImage(image) is not None
+    except ResourceNotFoundError:
+      return False
+
+  def GetCommonInstanceMetadata(self, key):
+    """Looks up a single project metadata value.
+
+    Args:
+      key: Metadata key name.
+
+    Returns:
+      Metadata value corresponding to the key, or None if it was not found.
+    """
+    projects_data = self.gce_client.projects().get(
+        project=self.project).execute()
+    metadata = projects_data['commonInstanceMetadata']
+    return _GetMetdataValue(metadata, key)
+
+  def SetCommonInstanceMetadata(self, key, value):
+    """Sets a single project metadata value.
+
+    Args:
+      key: Metadata key to be set.
+      value: New value, or None if the given key should be removed.
+    """
+    projects_data = self.gce_client.projects().get(
+        project=self.project).execute()
+    metadata = projects_data['commonInstanceMetadata']
+    _UpdateMetadataValue(metadata, key, value)
+    operation = self.gce_client.projects().setCommonInstanceMetadata(
+        project=self.project,
+        body=metadata).execute()
+    self._WaitForGlobalOperation(operation['name'])
+
+  def GetInstanceMetadata(self, instance, key):
+    """Looks up instance's metadata value.
+
+    Args:
+      instance: Name of the instance.
+      key: Metadata key name.
+
+    Returns:
+      Metadata value corresponding to the key, or None if it was not found.
+    """
+    instance_data = self.GetInstance(instance)
+    metadata = instance_data['metadata']
+    return self._GetMetdataValue(metadata, key)
+
+  def SetInstanceMetadata(self, instance, key, value):
+    """Sets a single instance metadata value.
+
+    Args:
+      instance: Name of the instance.
+      key: Metadata key to be set.
+      value: New value, or None if the given key should be removed.
+    """
+    instance_data = self.GetInstance(instance)
+    metadata = instance_data['metadata']
+    _UpdateMetadataValue(metadata, key, value)
+    operation = self.gce_client.instances().setMetadata(
+        project=self.project,
+        zone=self.zone,
+        instance=instance,
+        body=metadata).execute()
+    self._WaitForZoneOperation(operation['name'])
+
+  def _WaitForZoneOperation(self, operation, zone=None, timeout_sec=None,
+                            timeout_handler=None):
+    """Waits until a GCE ZoneOperation is finished or timed out.
+
+    Args:
+      operation: The GCE operation to wait for.
+      zone: The zone that |operation| belongs to.
+      timeout_sec: The maximum number of seconds to wait for.
+      timeout_handler: A callable to be executed when timeout happens.
+
+    Raises:
+      Error when timeout happens or the operation fails.
+    """
+    get_request = self.gce_client.zoneOperations().get(
+        project=self.project, zone=zone or self.zone, operation=operation)
+    self._WaitForOperation(operation, get_request, timeout_sec,
+                           timeout_handler=timeout_handler)
+
+  def _WaitForRegionOperation(self, operation, region, timeout_sec=None,
+                              timeout_handler=None):
+    """Waits until a GCE RegionOperation is finished or timed out.
+
+    Args:
+      operation: The GCE operation to wait for.
+      region: The region that |operation| belongs to.
+      timeout_sec: The maximum number of seconds to wait for.
+      timeout_handler: A callable to be executed when timeout happens.
+
+    Raises:
+      Error when timeout happens or the operation fails.
+    """
+    get_request = self.gce_client.regionOperations().get(
+        project=self.project, region=region or self.region, operation=operation)
+    self._WaitForOperation(operation, get_request, timeout_sec,
+                           timeout_handler=timeout_handler)
+
+  def _WaitForGlobalOperation(self, operation, timeout_sec=None,
+                              timeout_handler=None):
+    """Waits until a GCE GlobalOperation is finished or timed out.
+
+    Args:
+      operation: The GCE operation to wait for.
+      timeout_sec: The maximum number of seconds to wait for.
+      timeout_handler: A callable to be executed when timeout happens.
+
+    Raises:
+      Error when timeout happens or the operation fails.
+    """
+    get_request = self.gce_client.globalOperations().get(project=self.project,
+                                                         operation=operation)
+    self._WaitForOperation(operation, get_request, timeout_sec=timeout_sec,
+                           timeout_handler=timeout_handler)
+
+  def _WaitForOperation(self, operation, get_operation_request,
+                        timeout_sec=None, timeout_handler=None):
+    """Waits until timeout or the request gets a response with a 'DONE' status.
+
+    Args:
+      operation: The GCE operation to wait for.
+      get_operation_request:
+        The HTTP request to get the operation's status.
+        This request will be executed periodically until it returns a status
+        'DONE'.
+      timeout_sec: The maximum number of seconds to wait for.
+      timeout_handler: A callable to be executed when times out.
+
+    Raises:
+      Error when timeout happens or the operation fails.
+    """
+    def _IsDone():
+      result = get_operation_request.execute()
+      if result['status'] == 'DONE':
+        if 'error' in result:
+          raise Error(result['error'])
+        return True
+      return False
+
+    try:
+      timeout = timeout_sec or self.DEFAULT_TIMEOUT_SEC
+      logging.info('Waiting up to %d seconds for operation [%s] to complete...',
+                   timeout, operation)
+      timeout_util.WaitForReturnTrue(_IsDone, timeout, period=1)
+    except timeout_util.TimeoutError:
+      if timeout_handler:
+        timeout_handler()
+      raise Error('Timeout wating for operation [%s] to complete' % operation)
+
+  def _BuildRetriableRequest(self, num_retries, http, thread_safe=False,
+                             credentials=None, *args, **kwargs):
+    """Builds a request that will be automatically retried on server errors.
+
+    Args:
+      num_retries: The maximum number of times to retry until give up.
+      http: An httplib2.Http object that this request will be executed through.
+      thread_safe: Whether or not the request needs to be thread-safe.
+      credentials: Credentials to apply to the request.
+      *args: Optional positional arguments.
+      **kwargs: Optional keyword arguments.
+
+    Returns:
+      RetryOnServerErrorHttpRequest: A request that will automatically retried
+          on server errors.
+    """
+    if thread_safe:
+      # Create a new http object for every request.
+      http = credentials.authorize(httplib2.Http())
+    return RetryOnServerErrorHttpRequest(num_retries, http, *args, **kwargs)
diff --git a/au_test_harness/gce_au_worker.py b/au_test_harness/gce_au_worker.py
index ac606c9..e63d5e4 100644
--- a/au_test_harness/gce_au_worker.py
+++ b/au_test_harness/gce_au_worker.py
@@ -19,7 +19,7 @@
 
 from chromite.lib import cros_build_lib
 from chromite.lib import cros_logging as logging
-from chromite.lib import gce
+import gce  # TODO(crbug.com/1170006): Move back to chromite.
 from chromite.lib import gs
 from chromite.lib import path_util
 from crostestutils.au_test_harness import au_worker
diff --git a/au_test_harness/gce_au_worker_unittest.py b/au_test_harness/gce_au_worker_unittest.py
index 2bf7761..379dfa4 100755
--- a/au_test_harness/gce_au_worker_unittest.py
+++ b/au_test_harness/gce_au_worker_unittest.py
@@ -17,9 +17,9 @@
 sys.path.append(constants.CROS_PLATFORM_ROOT)
 sys.path.append(constants.SOURCE_ROOT)
 
-from chromite.lib.gce import GceContext
 from chromite.lib import cros_build_lib
 from chromite.lib import cros_test_lib
+import gce  # TODO(crbug.com/1170006): Move back to chromite.
 from chromite.lib import osutils
 from chromite.lib import path_util
 from crostestutils.au_test_harness.au_worker import AUWorker
@@ -55,8 +55,8 @@
     # Mock out model or class level methods.
     self.PatchObject(AUWorker, 'GetNextResultsPath', autospec=True,
                      return_value=('./log', './fail'))
-    self.PatchObject(GceContext, 'ForServiceAccountThreadSafe',
-                     spec=GceContext.ForServiceAccountThreadSafe)
+    self.PatchObject(gce.GceContext, 'ForServiceAccountThreadSafe',
+                     spec=gce.GceContext.ForServiceAccountThreadSafe)
 
   def testCreateInstance(self):
     """Verifies _CreateInstance creates an instance with |image_path|.