def test_get_cluster(self):
# Setup Expected Response
project_id_2 = "projectId2939242356"
cluster_name_2 = "clusterName2875867491"
cluster_uuid = "clusterUuid-1017854240"
expected_response = {
"project_id": project_id_2,
"cluster_name": cluster_name_2,
"cluster_uuid": cluster_uuid,
}
expected_response = clusters_pb2.Cluster(**expected_response)
# Mock the API response
channel = ChannelStub(responses=[expected_response])
patch = mock.patch("google.api_core.grpc_helpers.create_channel")
with patch as create_channel:
create_channel.return_value = channel
client = dataproc_v1beta2.ClusterControllerClient()
# Setup Request
project_id = "projectId-1969970175"
region = "region-934795532"
cluster_name = "clusterName-1018081872"
response = client.get_cluster(project_id, region, cluster_name)
assert expected_response == response
assert len(channel.requests) == 1
expected_request = clusters_pb2.GetClusterRequest(
project_id=project_id, region=region, cluster_name=cluster_name
)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
def test_create_cluster(self):
# Setup Expected Response
project_id_2 = 'projectId2939242356'
cluster_name = 'clusterName-1018081872'
cluster_uuid = 'clusterUuid-1017854240'
expected_response = {
'project_id': project_id_2,
'cluster_name': cluster_name,
'cluster_uuid': cluster_uuid
}
expected_response = clusters_pb2.Cluster(**expected_response)
operation = operations_pb2.Operation(
name='operations/test_create_cluster', done=True)
operation.response.Pack(expected_response)
# Mock the API response
channel = ChannelStub(responses=[operation])
client = dataproc_v1beta2.ClusterControllerClient(channel=channel)
# Setup Request
project_id = 'projectId-1969970175'
region = 'region-934795532'
cluster = {}
response = client.create_cluster(project_id, region, cluster)
result = response.result()
assert expected_response == result
assert len(channel.requests) == 1
expected_request = clusters_pb2.CreateClusterRequest(
project_id=project_id, region=region, cluster=cluster)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
def test_get_cluster(self):
# Setup Expected Response
project_id_2 = 'projectId2939242356'
cluster_name_2 = 'clusterName2875867491'
cluster_uuid = 'clusterUuid-1017854240'
expected_response = {
'project_id': project_id_2,
'cluster_name': cluster_name_2,
'cluster_uuid': cluster_uuid
}
expected_response = clusters_pb2.Cluster(**expected_response)
# Mock the API response
channel = ChannelStub(responses=[expected_response])
client = dataproc_v1beta2.ClusterControllerClient(channel=channel)
# Setup Request
project_id = 'projectId-1969970175'
region = 'region-934795532'
cluster_name = 'clusterName-1018081872'
response = client.get_cluster(project_id, region, cluster_name)
assert expected_response == response
assert len(channel.requests) == 1
expected_request = clusters_pb2.GetClusterRequest(
project_id=project_id, region=region, cluster_name=cluster_name)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
def create_cluster(cluster_json_path, cluster_name, region):
"""
This function takes an external JSON with cluster
configuration and creates the cluster.
:param project_id: project id of the service account
:param region: region where needs to be created
:return: returns the cluster client instance
"""
cluster_client = dataproc_v1beta2.ClusterControllerClient(
client_options={
'api_endpoint': '{}-dataproc.googleapis.com:443'.format(region)
})
# The json is created under the cluster config folder
with open(cluster_json_path, 'r') as f:
cluster_json_template = Template(f.read())
cluster_json = cluster_json_template.render(cluster_name=cluster_name)
cluster_config = json_format.Parse(cluster_json, clusters.Cluster())
project_id = cluster_config.project_id
operation = cluster_client.create_cluster(project_id, region,
cluster_config)
result = operation.result()
# Output a success message.
print('Cluster created successfully : {}'.format(result.cluster_name))
def test_update_cluster(self):
# Setup Expected Response
project_id_2 = "projectId2939242356"
cluster_name_2 = "clusterName2875867491"
cluster_uuid = "clusterUuid-1017854240"
expected_response = {
"project_id": project_id_2,
"cluster_name": cluster_name_2,
"cluster_uuid": cluster_uuid,
}
expected_response = clusters_pb2.Cluster(**expected_response)
operation = operations_pb2.Operation(
name="operations/test_update_cluster", done=True
)
operation.response.Pack(expected_response)
# Mock the API response
channel = ChannelStub(responses=[operation])
patch = mock.patch("google.api_core.grpc_helpers.create_channel")
with patch as create_channel:
create_channel.return_value = channel
client = dataproc_v1beta2.ClusterControllerClient()
# Setup Request
project_id = "projectId-1969970175"
region = "region-934795532"
cluster_name = "clusterName-1018081872"
cluster = {}
update_mask = {}
response = client.update_cluster(
project_id, region, cluster_name, cluster, update_mask
)
result = response.result()
assert expected_response == result
assert len(channel.requests) == 1
expected_request = clusters_pb2.UpdateClusterRequest(
project_id=project_id,
region=region,
cluster_name=cluster_name,
cluster=cluster,
update_mask=update_mask,
)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
async def test_poll_create(self):
expected_response = {
"status": {
"state": "CREATING"
}
}
expected_response = clusters_pb2.Cluster(**expected_response)
mock_client = mock.create_autospec(dataproc_v1beta2.ClusterControllerClient())
mock_client.get_cluster.return_value = expected_response
spawner = DataprocSpawner(hub=Hub(), dataproc=mock_client, user=MockUser(), _mock=True)
spawner.project = "test-poll-create"
assert spawner.project == "test-poll-create"
assert await spawner.poll() == None
