Skip to content

Commit c969b83

Browse files
authored
fix: add mains to samples (#3284)
Added mains to two samples: create_cluster and instantiate_inline_workflow_templates. Fixed their associated tests to accommodate this. Removed subprocess from quickstart/quickstart_test.py to fix [2873](#2873) fixes #2873
1 parent 05ca1c8 commit c969b83

File tree

4 files changed

+37
-17
lines changed

4 files changed

+37
-17
lines changed

dataproc/create_cluster.py

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,12 @@
1616

1717
# This sample walks a user through creating a Cloud Dataproc cluster using
1818
# the Python client library.
19+
#
20+
# This script can be run on its own:
21+
# python create_cluster.py ${PROJECT_ID} ${REGION} ${CLUSTER_NAME}
22+
23+
24+
import sys
1925

2026
# [START dataproc_create_cluster]
2127
from google.cloud import dataproc_v1 as dataproc
@@ -33,7 +39,7 @@ def create_cluster(project_id, region, cluster_name):
3339

3440
# Create a client with the endpoint set to the desired cluster region.
3541
cluster_client = dataproc.ClusterControllerClient(client_options={
36-
'api_endpoint': '{}-dataproc.googleapis.com:443'.format(region)
42+
'api_endpoint': f'{region}-dataproc.googleapis.com:443',
3743
})
3844

3945
# Create the cluster config.
@@ -57,5 +63,15 @@ def create_cluster(project_id, region, cluster_name):
5763
result = operation.result()
5864

5965
# Output a success message.
60-
print('Cluster created successfully: {}'.format(result.cluster_name))
66+
print(f'Cluster created successfully: {result.cluster_name}')
6167
# [END dataproc_create_cluster]
68+
69+
70+
if __name__ == "__main__":
71+
if len(sys.argv) < 4:
72+
sys.exit('python create_cluster.py project_id region cluster_name')
73+
74+
project_id = sys.argv[1]
75+
region = sys.argv[2]
76+
cluster_name = sys.argv[3]
77+
create_cluster(project_id, region, cluster_name)

dataproc/create_cluster_test.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ def teardown():
3131
yield
3232

3333
cluster_client = dataproc.ClusterControllerClient(client_options={
34-
'api_endpoint': '{}-dataproc.googleapis.com:443'.format(REGION)
34+
'api_endpoint': f'{REGION}-dataproc.googleapis.com:443'
3535
})
3636
# Client library function
3737
operation = cluster_client.delete_cluster(PROJECT_ID, REGION, CLUSTER_NAME)

dataproc/instantiate_inline_workflow_template.py

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,11 @@
1616
# workflow for Cloud Dataproc using the Python client library.
1717
#
1818
# This script can be run on its own:
19-
# python workflows.py ${PROJECT_ID} ${REGION}
19+
# python instantiate_inline_workflow_template.py ${PROJECT_ID} ${REGION}
20+
2021

2122
import sys
23+
2224
# [START dataproc_instantiate_inline_workflow_template]
2325
from google.cloud import dataproc_v1 as dataproc
2426

@@ -35,7 +37,8 @@ def instantiate_inline_workflow_template(project_id, region):
3537
# Create a client with the endpoint set to the desired region.
3638
workflow_template_client = dataproc.WorkflowTemplateServiceClient(
3739
client_options={
38-
'api_endpoint': '{}-dataproc.googleapis.com:443'.format(region)}
40+
'api_endpoint': f'{region}-dataproc.googleapis.com:443'
41+
}
3942
)
4043

4144
parent = workflow_template_client.region_path(project_id, region)
@@ -91,8 +94,14 @@ def instantiate_inline_workflow_template(project_id, region):
9194

9295
# Output a success message.
9396
print('Workflow ran successfully.')
94-
# [END dataproc_instantiate_inline_workflow_template]
97+
# [END dataproc_instantiate_inline_workflow_template]
9598

9699

97100
if __name__ == "__main__":
98-
instantiate_inline_workflow_template(sys.argv[1], sys.argv[2])
101+
if len(sys.argv) < 3:
102+
sys.exit('python instantiate_inline_workflow_template.py '
103+
+ 'project_id region')
104+
105+
project_id = sys.argv[1]
106+
region = sys.argv[2]
107+
instantiate_inline_workflow_template(project_id, region)

dataproc/quickstart/quickstart_test.py

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,12 @@
1515
import os
1616
import uuid
1717
import pytest
18-
import subprocess
1918

2019
from google.cloud import dataproc_v1 as dataproc
2120
from google.cloud import storage
2221

22+
import quickstart
23+
2324

2425
PROJECT_ID = os.environ['GCLOUD_PROJECT']
2526
REGION = 'us-central1'
@@ -60,15 +61,9 @@ def setup_teardown():
6061
bucket.delete()
6162

6263

63-
def test_quickstart():
64-
command = [
65-
'python', 'quickstart/quickstart.py',
66-
'--project_id', PROJECT_ID,
67-
'--region', REGION,
68-
'--cluster_name', CLUSTER_NAME,
69-
'--job_file_path', JOB_FILE_PATH
70-
]
71-
out = subprocess.check_output(command).decode("utf-8")
64+
def test_quickstart(capsys):
65+
quickstart.quickstart(PROJECT_ID, REGION, CLUSTER_NAME, JOB_FILE_PATH)
66+
out, _ = capsys.readouterr()
7267

7368
assert 'Cluster created successfully' in out
7469
assert 'Submitted job' in out

0 commit comments

Comments
 (0)