sdevani commented on a change in pull request #5251: [AIRFLOW-4135] Add Google 
Cloud Build operator and hook
URL: https://github.com/apache/airflow/pull/5251#discussion_r283459152
 
 

 ##########
 File path: tests/contrib/operators/test_gcp_cloud_build_operator.py
 ##########
 @@ -0,0 +1,194 @@
+from copy import deepcopy
+from unittest import TestCase
+
+import six
+
+from parameterized import parameterized
+from airflow import AirflowException
+from airflow.contrib.operators.gcp_cloud_build_operator import BuildProcessor, 
CloudBuildCreateBuildOperator
+
+try:
+    # noinspection PyProtectedMember
+    from unittest import mock
+except ImportError:
+    try:
+        import mock
+    except ImportError:
+        mock = None
+
+TEST_CREATE_BODY = {
+    "source": {"storageSource": {"bucket": "cloud-build-examples", "object": 
"node-docker-example.tar.gz"}},
+    "steps": [
+        {"name": "gcr.io/cloud-builders/docker", "args": ["build", "-t", 
"gcr.io/$PROJECT_ID/my-image", "."]}
+    ],
+    "images": ["gcr.io/$PROJECT_ID/my-image"],
+}
+TEST_PROJECT_ID = "example-id"
+
+
+class BuildProcessorTestCase(TestCase):
+    def test_verify_source(self):
+        with six.assertRaisesRegex(self, AirflowException, "The source could 
not be determined."):
+            BuildProcessor(body={"source": {"storageSource": {}, "repoSource": 
{}}}).process_body()
+
+    @parameterized.expand(
+        [
+            (
+                
"https://source.developers.google.com/p/airflow-project/r/airflow-repo";,
+                {"projectId": "airflow-project", "repoName": "airflow-repo", 
"branchName": "master"},
+            ),
+            (
+                
"https://source.developers.google.com/p/airflow-project/r/airflow-repo#branch-name";,
+                {"projectId": "airflow-project", "repoName": "airflow-repo", 
"branchName": "branch-name"},
+            ),
+            (
+                
"https://source.developers.google.com/p/airflow-project/r/airflow-repo#feature/branch";,
+                {"projectId": "airflow-project", "repoName": "airflow-repo", 
"branchName": "feature/branch"},
+            ),
+        ]
+    )
+    def test_convert_repo_url_to_dict_valid(self, url, expected_dict):
+        body = {"source": {"repoSource": url}}
+        body = BuildProcessor(body=body).process_body()
+        self.assertEquals(body["source"]["repoSource"], expected_dict)
+
+    @parameterized.expand(
+        [
+            
("https://source.e.com/p/airflow-project/r/airflow-repo#branch-name";,),
+            
("httpXs://source.developers.google.com/p/airflow-project/r/airflow-repo",),
+            
("://source.developers.google.com/p/airflow-project/r/airflow-repo",),
+            
("://source.developers.google.com/p/airflow-project/rXXXX/airflow-repo",),
+            
("://source.developers.google.com/pXXX/airflow-project/r/airflow-repo",),
+        ]
+    )
+    def test_convert_repo_url_to_storage_dict_invalid(self, url):
+        body = {"source": {"repoSource": url}}
+        with six.assertRaisesRegex(self, AirflowException, "Invalid URL."):
+            BuildProcessor(body=body).process_body()
+
+    @parameterized.expand(
+        [
+            (
+                "gs://bucket-name/airflow-object.tar.gz",
+                {"bucket": "bucket-name", "object": "airflow-object.tar.gz"},
+            ),
+            (
+                "gs://bucket-name/airflow-object.tar.gz#1231231",
+                {"bucket": "bucket-name", "object": "airflow-object.tar.gz", 
"generation": "1231231"},
+            ),
+        ]
+    )
+    def test_convert_storage_url_to_dict_valid(self, url, expected_dict):
+        body = {"source": {"storageSource": url}}
+        body = BuildProcessor(body=body).process_body()
+        self.assertEquals(body["source"]["storageSource"], expected_dict)
+
+    @parameterized.expand(
+        [("///object",), ("gsXXa:///object",), ("gs://bucket-name/",), 
("gs://bucket-name",)]
+    )
+    def test_convert_storage_url_to_dict_invalid(self, url):
+        body = {"source": {"storageSource": url}}
+        with six.assertRaisesRegex(self, AirflowException, "Invalid URL."):
+            BuildProcessor(body=body).process_body()
+
+    @parameterized.expand([("storageSource",), ("repoSource",)])
+    def test_do_nothing(self, source_key):
+        body = {"source": {source_key: {}}}
+        expected_body = deepcopy(body)
+
+        BuildProcessor(body=body).process_body()
+        self.assertEquals(body, expected_body)
+
+
+class GcpCloudBuildCreateBuildOperatorTestCase(TestCase):
+    @mock.patch(
+        "airflow.contrib.operators.gcp_cloud_build_operator.CloudBuildHook",
+        **{"return_value.create_build.return_value": TEST_CREATE_BODY}
+    )
+    def test_minimal_green_path(self, hock_mock):
+        op = CloudBuildCreateBuildOperator(
+            body=TEST_CREATE_BODY, project_id=TEST_PROJECT_ID, 
task_id="task-id"
+        )
+        r = op.execute({})
+        self.assertIs(r, TEST_CREATE_BODY)
+
+    @parameterized.expand([({},), (None,)])
+    def test_missing_input(self, body):
+        with six.assertRaisesRegex(self, AirflowException, "The required 
parameter 'body' is missing"):
+            CloudBuildCreateBuildOperator(body=body, 
project_id=TEST_PROJECT_ID, task_id="task-id")
+
+    @mock.patch(
+        "airflow.contrib.operators.gcp_cloud_build_operator.CloudBuildHook",
+        **{"return_value.create_build.return_value": TEST_CREATE_BODY}
+    )
+    def test_storage_source_replace(self, hock_mock):
+        current_body = {
+            # [START howto_operator_gcp_cloud_build_source_gcs_url]
+            "source": {"storageSource": "gs://bucket-name/object-name.tar.gz"},
+            # [END howto_operator_gcp_cloud_build_source_gcs_url]
+            "steps": [
+                {
+                    "name": "gcr.io/cloud-builders/docker",
+                    "args": ["build", "-t", "gcr.io/$PROJECT_ID/docker-image", 
"."],
+                }
+            ],
+            "images": ["gcr.io/$PROJECT_ID/docker-image"],
+        }
+
+        op = CloudBuildCreateBuildOperator(body=current_body, 
project_id=TEST_PROJECT_ID, task_id="task-id")
+        op.execute({})
+
+        expected_result = {
+            # [START howto_operator_gcp_cloud_build_source_gcs_dict]
+            "source": {"storageSource": {"bucket": "bucket-name", "object": 
"object-name.tar.gz"}},
+            # [END howto_operator_gcp_cloud_build_source_gcs_dict]
+            "steps": [
+                {
+                    "name": "gcr.io/cloud-builders/docker",
+                    "args": ["build", "-t", "gcr.io/$PROJECT_ID/docker-image", 
"."],
+                }
+            ],
+            "images": ["gcr.io/$PROJECT_ID/docker-image"],
+        }
+        hock_mock.create_build(body=expected_result, 
project_id=TEST_PROJECT_ID)
 
 Review comment:
   There is no assertion in this test

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
[email protected]


With regards,
Apache Git Services

Reply via email to