37
37
38
38
# The absolute path of the current file. This will locate the model_path when
39
39
# 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__ ) ))
41
41
MODEL_PATH = os .path .join (ROOT_DIR , 'model_path' )
42
+
43
+ IMAGE_FILE_PATH = os .path .join (os .path .dirname (__file__ ), 'test.jpg' )
42
44
# 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'
45
46
# 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'
48
48
# Container Name.
49
49
NAME = 'AutomlVisionEdgeContainerPredictTest'
50
50
# Port Number.
@@ -56,8 +56,8 @@ def edge_container_predict_server_port():
56
56
# set up
57
57
# Pull the CPU docker.
58
58
subprocess .check_output (['docker' , 'pull' , CPU_DOCKER_GCS_PATH ])
59
- # Get the sample saved model.
60
59
60
+ # Get the sample saved model.
61
61
if not os .path .exists (MODEL_PATH ):
62
62
os .mkdir (MODEL_PATH )
63
63
subprocess .check_output (
@@ -79,18 +79,16 @@ def edge_container_predict_server_port():
79
79
# Remove the docker image.
80
80
subprocess .check_output (['docker' , 'rmi' , CPU_DOCKER_GCS_PATH ])
81
81
82
- # TODO(dizcology): Enable tests in future.
83
- @pytest .mark .skip (reason = 'skipping to avoid running docker in docker' )
82
+
84
83
def test_edge_container_predict (capsys , edge_container_predict_server_port ):
85
- image_file_path = 'test.jpg'
86
84
# If you send requests with one image each time, the key value does not
87
85
# matter. If you send requests with multiple images, please used different
88
86
# keys to indicated different images, which can make sure that the
89
87
# responses corresponding to the given image.
90
88
image_key = '1'
91
89
# Send a request.
92
90
response = predict .container_predict (
93
- image_file_path , image_key , PORT_NUMBER )
91
+ IMAGE_FILE_PATH , image_key , PORT_NUMBER )
94
92
# Verify the response.
95
93
assert 'predictions' in response
96
94
assert 'key' in response ['predictions' ][0 ]
0 commit comments