update test

dizcology · dizcology · commit 89478ee48991 · 2019-06-05T10:36:46.000-07:00
diff --git a/vision/automl/edge_container_predict/automl_vision_edge_container_predict_test.py b/vision/automl/edge_container_predict/automl_vision_edge_container_predict_test.py
@@ -37,14 +37,12 @@
 
 # The absolute path of the current file. This will locate the model_path when
 # run docker containers.
-ROOT_DIR = os.path.abspath(os.path.dirname(__file__))
+ROOT_DIR = os.environ.get('KOKORO_ROOT', os.path.abspath(os.path.dirname(__file__)))
 MODEL_PATH = os.path.join(ROOT_DIR, 'model_path')
 # The cpu docker gcs path is from 'Edge container tutorial'.
-DOCKER_GCS_DIR = 'gcr.io/automl-vision-ondevice/'
-CPU_DOCKER_GCS_PATH = DOCKER_GCS_DIR + 'gcloud-container-1.12.0:latest'
+CPU_DOCKER_GCS_PATH = 'gcr.io/automl-vision-ondevice/gcloud-container-1.12.0:latest'
 # The path of a sample saved model.
-MODEL_GCS_DIR = 'gs://cloud-samples-data/vision/edge_container_predict/'
-SAMPLE_SAVED_MODEL = MODEL_GCS_DIR + 'saved_model.pb'
+SAMPLE_SAVED_MODEL = 'gs://cloud-samples-data/vision/edge_container_predict/saved_model.pb'
 # Container Name.
 NAME = 'AutomlVisionEdgeContainerPredictTest'
 # Port Number.
@@ -56,8 +54,8 @@ def edge_container_predict_server_port():
     # set up
     # Pull the CPU docker.
     subprocess.check_output(['docker', 'pull', CPU_DOCKER_GCS_PATH])
-    # Get the sample saved model.
 
+    # Get the sample saved model.
     if not os.path.exists(MODEL_PATH):
         os.mkdir(MODEL_PATH)
     subprocess.check_output(
@@ -79,8 +77,7 @@ def edge_container_predict_server_port():
     # Remove the docker image.
     subprocess.check_output(['docker', 'rmi', CPU_DOCKER_GCS_PATH])
 
-# TODO(dizcology): Enable tests in future.
-@pytest.mark.skip(reason='skipping to avoid running docker in docker')
+
 def test_edge_container_predict(capsys, edge_container_predict_server_port):
     image_file_path = 'test.jpg'
     # If you send requests with one image each time, the key value does not
