feat: batch create with custom network sample (#4872)

* feat: batch create with custom network sample

* added buffer check

batch/create_with_custom_network.go

@@ -0,0 +1,117 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package snippets
+
+// [START batch_create_custom_network]
+import (
+	"context"
+	"fmt"
+	"io"
+
+	batch "cloud.google.com/go/batch/apiv1"
+	"cloud.google.com/go/batch/apiv1/batchpb"
+	durationpb "google.golang.org/protobuf/types/known/durationpb"
+)
+
+// Creates and runs a job with custom network
+func createJobWithCustomNetwork(w io.Writer, projectID, region, jobName, networkName, subnetworkName string) (*batchpb.Job, error) {
+	ctx := context.Background()
+	batchClient, err := batch.NewClient(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("batchClient error: %w", err)
+	}
+	defer batchClient.Close()
+
+	runn := &batchpb.Runnable{
+		Executable: &batchpb.Runnable_Script_{
+			Script: &batchpb.Runnable_Script{
+				Command: &batchpb.Runnable_Script_Text{
+					Text: "echo Hello world from script 1 for task ${BATCH_TASK_INDEX}",
+				},
+			},
+		},
+	}
+
+	taskSpec := &batchpb.TaskSpec{
+		ComputeResource: &batchpb.ComputeResource{
+			// CpuMilli is milliseconds per cpu-second. This means the task requires 2 whole CPUs.
+			CpuMilli:  2000,
+			MemoryMib: 16,
+		},
+		MaxRunDuration: &durationpb.Duration{
+			Seconds: 3600,
+		},
+		MaxRetryCount: 2,
+		Runnables:     []*batchpb.Runnable{runn},
+	}
+
+	taskGroups := []*batchpb.TaskGroup{
+		{
+			TaskCount: 4,
+			TaskSpec:  taskSpec,
+		},
+	}
+
+	// Policies are used to define on what kind of virtual machines the tasks will run on.
+	// In this case, we tell the system to use "e2-standard-4" machine type.
+	// Read more about machine types here: https://cloud.google.com/compute/docs/machine-types
+	allocationPolicy := &batchpb.AllocationPolicy{
+		Instances: []*batchpb.AllocationPolicy_InstancePolicyOrTemplate{{
+			PolicyTemplate: &batchpb.AllocationPolicy_InstancePolicyOrTemplate_Policy{
+				Policy: &batchpb.AllocationPolicy_InstancePolicy{
+					MachineType: "e2-standard-4",
+				},
+			},
+		}},
+		Network: &batchpb.AllocationPolicy_NetworkPolicy{
+			NetworkInterfaces: []*batchpb.AllocationPolicy_NetworkInterface{
+				{
+					// Set the network to the specified network name within the project
+					Network: fmt.Sprintf("projects/%s/global/networks/%s", projectID, networkName),
+					// Set the subnetwork to the specified subnetwork within the region
+					Subnetwork: fmt.Sprintf("projects/%s/regions/%s/subnetworks/%s", projectID, region, subnetworkName),
+				},
+			},
+		},
+	}
+
+	// We use Cloud Logging as it's an out of the box available option
+	logsPolicy := &batchpb.LogsPolicy{
+		Destination: batchpb.LogsPolicy_CLOUD_LOGGING,
+	}
+
+	job := &batchpb.Job{
+		Name:             jobName,
+		TaskGroups:       taskGroups,
+		AllocationPolicy: allocationPolicy,
+		LogsPolicy:       logsPolicy,
+	}
+
+	request := &batchpb.CreateJobRequest{
+		Parent: fmt.Sprintf("projects/%s/locations/%s", projectID, region),
+		JobId:  jobName,
+		Job:    job,
+	}
+
+	created_job, err := batchClient.CreateJob(ctx, request)
+	if err != nil {
+		return nil, fmt.Errorf("unable to create job: %w", err)
+	}
+
+	fmt.Fprintf(w, "Job created: %v\n", created_job)
+	return created_job, nil
+}
+
+// [END batch_create_custom_network]

batch/create_with_custom_network_test.go

@@ -0,0 +1,58 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package snippets
+
+import (
+	"bytes"
+	"fmt"
+	"math/rand"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/GoogleCloudPlatform/golang-samples/internal/testutil"
+)
+
+func TestCreateJobWithCustomNetwork(t *testing.T) {
+	var r *rand.Rand = rand.New(
+		rand.NewSource(time.Now().UnixNano()))
+	tc := testutil.SystemTest(t)
+	jobName := fmt.Sprintf("test-job-go-%v-%v", time.Now().Format("2006-01-02"), r.Int())
+	region := "us-central1"
+	networkName, subnetworkName := "default", "default"
+
+	buf := &bytes.Buffer{}
+
+	job, err := createJobWithCustomNetwork(buf, tc.ProjectID, region, jobName, networkName, subnetworkName)
+
+	if err != nil {
+		t.Errorf("createJobWithCustomNetwork got err: %v", err)
+	}
+	if got := buf.String(); !strings.Contains(got, "Job created") {
+		t.Errorf("createJobWithCustomNetwork got %q, expected %q", got, "Job created")
+	}
+
+	expectedNetwork := fmt.Sprintf("projects/%s/global/networks/%s", tc.ProjectID, networkName)
+	expectedSubnetwork := fmt.Sprintf("projects/%s/regions/%s/subnetworks/%s", tc.ProjectID, region, subnetworkName)
+
+	interfaces := job.GetAllocationPolicy().GetNetwork().GetNetworkInterfaces()
+	if interfaces[0].GetNetwork() != expectedNetwork || interfaces[0].GetSubnetwork() != expectedSubnetwork {
+		t.Errorf("Network wasn't set")
+	}
+
+	if err := deleteJob(buf, tc.ProjectID, region, jobName); err != nil {
+		t.Errorf("deleteJob got err: %v", err)
+	}
+}