Video Intelligence: Enrich VPCSC tests.

videointelligence/tests/system.py

@@ -20,16 +20,27 @@
 import time
 import unittest
 
+from google.auth.transport import requests as goog_auth_requests
 from google.cloud import videointelligence
 from google.cloud.videointelligence_v1 import enums
+from google.oauth2 import service_account
 
-PROJECT_NUMBER = os.environ.get("PROJECT_NUMBER")
-OUTSIDE_PROJECT_API_KEY = os.environ.get(
-    "GOOGLE_CLOUD_TESTS_VPCSC_OUTSIDE_PERIMETER_PROJECT_API_KEY"
-)
+CLOUD_PLATFORM_SCOPE = 'https://www.googleapis.com/auth/cloud-platform'
+CREDENTIALS_FILE = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
 OUTSIDE_IP = os.environ.get("GOOGLE_CLOUD_TESTS_VPCSC_OUTSIDE_IP")
 INSIDE_IP = os.environ.get("GOOGLE_CLOUD_TESTS_VPCSC_INSIDE_IP")
-
+IS_INSIDE_VPCSC = os.environ.get("GOOGLE_CLOUD_TESTS_IN_VPCSC")
+
+def get_access_token():
+  """Returns an access token.
+
+  Generates access tokens using the provided service account key file.
+  """
+  creds = service_account.Credentials.from_service_account_file(
+      CREDENTIALS_FILE, scopes=[CLOUD_PLATFORM_SCOPE])
+  with requests.Session() as session:
+    creds.refresh(goog_auth_requests.Request(session=session))
+  return creds.token
 
 class VideoIntelligenceSystemTestBase(unittest.TestCase):
     client = None
@@ -54,7 +65,7 @@ def test_annotate_video(self):
 
         # Wait for the operation to complete.
         # Long timeout value warranted due to https://github.com/grpc/grpc/issues/19173
-        lro_timeout_seconds = 180
+        lro_timeout_seconds = 300
         start_time = time.time()
         cnt = 0
         while not response.done() and (time.time() - start_time) < lro_timeout_seconds:
@@ -73,61 +84,72 @@ def test_annotate_video(self):
 
 
 @unittest.skipUnless(
-    OUTSIDE_PROJECT_API_KEY,
-    "GOOGLE_CLOUD_TESTS_VPCSC_OUTSIDE_PERIMETER_PROJECT_API_KEY not set in environment.",
+    CREDENTIALS_FILE,
+    "GOOGLE_APPLICATION_CREDENTIALS not set in environment.",
 )
 class TestVideoIntelligenceClientVpcSc(VideoIntelligenceSystemTestBase):
     # Tests to verify VideoIntelligence service requests blocked when trying to access resources outside of a secure perimeter.
     def setUp(self):
         VideoIntelligenceSystemTestBase.setUp(self)
         # api-endpoint
-        self.url = "https://videointelligence.googleapis.com/v1/videos:annotate?key={}".format(
-            OUTSIDE_PROJECT_API_KEY
-        )
+        self.url = "https://videointelligence.googleapis.com/v1/videos:annotate"
         self.body = {
             "input_uri": self.input_uri,
             "features": ["LABEL_DETECTION"],
             "location_id": "us-west1",
         }
 
-    @unittest.skipUnless(PROJECT_NUMBER, "PROJECT_NUMBER not set in environment.")
     @unittest.skipUnless(
         OUTSIDE_IP, "GOOGLE_CLOUD_TESTS_VPCSC_OUTSIDE_IP not set in environment."
     )
+    @unittest.skipUnless(
+        IS_INSIDE_VPCSC, "GOOGLE_CLOUD_TESTS_IN_VPCSC not set in environment."
+    )
     def test_outside_ip_address_blocked(self):
         headers = {
+            "Authorization": "Bearer " + get_access_token(),
             "Content-Type": "application/json",
             "X-User-IP": OUTSIDE_IP,
-            "X-Google-GFE-Cloud-Client-Network-Project-Number": PROJECT_NUMBER,
         }
         r = requests.post(url=self.url, data=json.dumps(self.body), headers=headers)
-        outside_project_operation = json.loads(r.text)
-        print(outside_project_operation)
+        operation = json.loads(r.text)
+        print(operation)
+
+        get_op_url = "https://videointelligence.googleapis.com/v1/" + operation["name"]
+        get_op = requests.get(url=get_op_url, headers=headers)
+        get_op_resp = json.loads(get_op.text)
+        print(get_op_resp)
         # Assert it returns permission denied from VPC SC
-        self.assertEqual(outside_project_operation["error"]["code"], 403)
+        self.assertEqual(get_op_resp["error"]["code"], 403)
         self.assertEqual(
-            outside_project_operation["error"]["status"], "PERMISSION_DENIED"
+            get_op_resp["error"]["status"], "PERMISSION_DENIED"
         )
         self.assertEqual(
-            outside_project_operation["error"]["details"][0]["violations"][0]["type"],
+            get_op_resp["error"]["details"][0]["violations"][0]["type"],
             "VPC_SERVICE_CONTROLS",
         )
         self.assertEqual(
-            outside_project_operation["error"]["message"],
+            get_op_resp["error"]["message"],
             "Request is prohibited by organization's policy",
         )
 
-    @unittest.skipUnless(PROJECT_NUMBER, "PROJECT_NUMBER not set in environment.")
     @unittest.skipUnless(
         INSIDE_IP, "GOOGLE_CLOUD_TESTS_VPCSC_INSIDE_IP not set in environment."
     )
     def test_inside_ip_address_allowed(self):
         headers = {
+            "Authorization": "Bearer " + get_access_token(),
             "Content-Type": "application/json",
             "X-User-IP": INSIDE_IP,
-            "X-Google-GFE-Cloud-Client-Network-Project-Number": PROJECT_NUMBER,
         }
         r = requests.post(url=self.url, data=json.dumps(self.body), headers=headers)
         operation = json.loads(r.text)
-        # Assert it returns non-empty operation name.
-        self.assertNotEqual(operation["name"], "")
+        print(operation)
+
+        get_op_url = "https://videointelligence.googleapis.com/v1/" + operation["name"]
+        get_op = requests.get(url=get_op_url, headers=headers)
+        get_op_resp = json.loads(get_op.text)
+        print(get_op_resp)
+        # Assert that we do not get an error.
+        self.assertEqual(get_op_resp["name"], operation["name"])
+