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