3737
3838# The absolute path of the current file. This will locate the model_path when
3939# run docker containers.
40- ROOT_DIR = os .path .abspath (os .path .dirname (__file__ ))
40+ ROOT_DIR = os .environ . get ( 'KOKORO_ROOT' , os . path .abspath (os .path .dirname (__file__ ) ))
4141MODEL_PATH = os .path .join (ROOT_DIR , 'model_path' )
42+
43+ IMAGE_FILE_PATH = os .path .join (os .path .dirname (__file__ ), 'test.jpg' )
4244# The cpu docker gcs path is from 'Edge container tutorial'.
43- DOCKER_GCS_DIR = 'gcr.io/automl-vision-ondevice/'
44- CPU_DOCKER_GCS_PATH = DOCKER_GCS_DIR + 'gcloud-container-1.12.0:latest'
45+ CPU_DOCKER_GCS_PATH = 'gcr.io/automl-vision-ondevice/gcloud-container-1.12.0:latest'
4546# The path of a sample saved model.
46- MODEL_GCS_DIR = 'gs://cloud-samples-data/vision/edge_container_predict/'
47- SAMPLE_SAVED_MODEL = MODEL_GCS_DIR + 'saved_model.pb'
47+ SAMPLE_SAVED_MODEL = 'gs://cloud-samples-data/vision/edge_container_predict/saved_model.pb'
4848# Container Name.
4949NAME = 'AutomlVisionEdgeContainerPredictTest'
5050# Port Number.
@@ -56,8 +56,8 @@ def edge_container_predict_server_port():
5656 # set up
5757 # Pull the CPU docker.
5858 subprocess .check_output (['docker' , 'pull' , CPU_DOCKER_GCS_PATH ])
59- # Get the sample saved model.
6059
60+ # Get the sample saved model.
6161 if not os .path .exists (MODEL_PATH ):
6262 os .mkdir (MODEL_PATH )
6363 subprocess .check_output (
@@ -79,18 +79,16 @@ def edge_container_predict_server_port():
7979 # Remove the docker image.
8080 subprocess .check_output (['docker' , 'rmi' , CPU_DOCKER_GCS_PATH ])
8181
82- # TODO(dizcology): Enable tests in future.
83- @pytest .mark .skip (reason = 'skipping to avoid running docker in docker' )
82+
8483def test_edge_container_predict (capsys , edge_container_predict_server_port ):
85- image_file_path = 'test.jpg'
8684 # If you send requests with one image each time, the key value does not
8785 # matter. If you send requests with multiple images, please used different
8886 # keys to indicated different images, which can make sure that the
8987 # responses corresponding to the given image.
9088 image_key = '1'
9189 # Send a request.
9290 response = predict .container_predict (
93- image_file_path , image_key , PORT_NUMBER )
91+ IMAGE_FILE_PATH , image_key , PORT_NUMBER )
9492 # Verify the response.
9593 assert 'predictions' in response
9694 assert 'key' in response ['predictions' ][0 ]
0 commit comments