diff --git a/codelabs/flex_and_vision/main.py b/codelabs/flex_and_vision/main.py
index 3425b120166..573f6d555bf 100644
--- a/codelabs/flex_and_vision/main.py
+++ b/codelabs/flex_and_vision/main.py
@@ -63,22 +63,24 @@ def upload_photo():
     blob.make_public()
 
     # Create a Cloud Vision client.
-    vision_client = vision.Client()
+    vision_client = vision.ImageAnnotatorClient()
 
     # Use the Cloud Vision client to detect a face for our image.
     source_uri = 'gs://{}/{}'.format(CLOUD_STORAGE_BUCKET, blob.name)
-    image = vision_client.image(source_uri=source_uri)
-    faces = image.detect_faces(limit=1)
+    image = vision.types.Image(
+        source=vision.types.ImageSource(gcs_image_uri=source_uri))
+    faces = vision_client.face_detection(image).face_annotations
 
     # If a face is detected, save to Datastore the likelihood that the face
     # displays 'joy,' as determined by Google's Machine Learning algorithm.
     if len(faces) > 0:
         face = faces[0]
 
-        # Convert the face.emotions.joy enum type to a string, which will be
-        # something like 'Likelihood.VERY_LIKELY'. Parse that string by the
-        # period to extract only the 'VERY_LIKELY' portion.
-        face_joy = str(face.emotions.joy).split('.')[1]
+        # Convert the likelihood string.
+        likelihoods = [
+            'Unknown', 'Very Unlikely', 'Unlikely', 'Possible', 'Likely',
+            'Very Likely']
+        face_joy = likelihoods[face.joy_likelihood]
     else:
         face_joy = 'Unknown'
 
diff --git a/codelabs/flex_and_vision/main_test.py b/codelabs/flex_and_vision/main_test.py
new file mode 100644
index 00000000000..c694656a6c3
--- /dev/null
+++ b/codelabs/flex_and_vision/main_test.py
@@ -0,0 +1,48 @@
+# Copyright 2017 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import requests
+import six
+
+import main
+
+TEST_PHOTO_URL = (
+    'https://upload.wikimedia.org/wikipedia/commons/5/5e/'
+    'John_F._Kennedy%2C_White_House_photo_portrait%2C_looking_up.jpg')
+
+
+@pytest.fixture
+def app():
+    main.app.testing = True
+    client = main.app.test_client()
+    return client
+
+
+def test_index(app):
+    r = app.get('/')
+    assert r.status_code == 200
+
+
+def test_upload_photo(app):
+    test_photo_data = requests.get(TEST_PHOTO_URL).content
+
+    r = app.post(
+        '/upload_photo',
+        data={
+            'file': (six.BytesIO(test_photo_data), 'flex_and_vision.jpg')
+        }
+    )
+
+    assert r.status_code == 302