Streaming go sample test

streaming/streaming.go

@@ -583,34 +583,58 @@
 	return nil
 }
 
-// LaunchDataflowJob populates the parameters from the streaming config and triggers a Dataflow job.
-func LaunchDataflowJob(ctx context.Context, targetProfile profiles.TargetProfile, streamingCfg StreamingCfg, conv *internal.Conv) (internal.DataflowOutput, error) {
-	project, instance, dbName, _ := targetProfile.GetResourceIds(ctx, time.Now(), "", nil)
-	dataflowCfg := streamingCfg.DataflowCfg
-	datastreamCfg := streamingCfg.DatastreamCfg
+type DataflowAccessor interface {
+	LaunchFlexTemplate(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error)
+}
 
-	// Rate limit this function to match DataFlow createJob Quota.
-	DATA_FLOW_RL.Take()
+type DataflowAccessorImpl struct{}
 
-	fmt.Println("Launching dataflow job ", dataflowCfg.JobName, " in ", project, "-", dataflowCfg.Location)
+func NewDataflowAccessor() DataflowAccessor {
+	return DataflowAccessorImpl{}
+}
 
-	c, err := dataflow.NewFlexTemplatesClient(ctx)
+func (dfA DataflowAccessorImpl) LaunchFlexTemplate(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error) {
+	fmt.Println("Created flex template client...")
+	dfClient, err := dataflow.NewFlexTemplatesClient(ctx)
+	defer dfClient.Close()
 	if err != nil {
-		return internal.DataflowOutput{}, fmt.Errorf("could not create flex template client: %v", err)
+		return nil, fmt.Errorf("could not create flex template client: %v", err)
 	}
-	defer c.Close()
-	fmt.Println("Created flex template client...")
+	return dfClient.LaunchFlexTemplate(ctx, req)
+}
 
-	//Creating datastream client to fetch the gcs bucket using target profile.
+type DatastreamAccessor interface {
+	GetConnectionProfile(ctx context.Context, req *datastreampb.GetConnectionProfileRequest, opts ...gax.CallOption) (*datastreampb.ConnectionProfile, error)
+}
+
+type DatastreamAccessorImpl struct{}
+
+func NewDatastreamAccessor() DatastreamAccessor {
+	return DatastreamAccessorImpl{}
+}
+
+func (dsA DatastreamAccessorImpl) GetConnectionProfile(ctx context.Context, req *datastreampb.GetConnectionProfileRequest, opts ...gax.CallOption) (*datastreampb.ConnectionProfile, error) {
 	dsClient, err := datastream.NewClient(ctx)
+	defer dsClient.Close()
 	if err != nil {
-		return internal.DataflowOutput{}, fmt.Errorf("datastream client can not be created: %v", err)
+		return nil, fmt.Errorf("datastream client can not be created: %v", err)
 	}
-	defer dsClient.Close()
+	return dsClient.GetConnectionProfile(ctx, req)
+}
 
+// LaunchDataflowJob populates the parameters from the streaming config and triggers a Dataflow job.
+func LaunchDataflowJob(ctx context.Context, targetProfile profiles.TargetProfile, streamingCfg StreamingCfg, conv *internal.Conv, dataflowAccessor DataflowAccessor, datastreamAccessor DatastreamAccessor) (internal.DataflowOutput, error) {
+	project, instance, dbName, _ := targetProfile.GetResourceIds(ctx, time.Now(), "", nil)
+	dataflowCfg := streamingCfg.DataflowCfg
+	datastreamCfg := streamingCfg.DatastreamCfg
+
+	// Rate limit this function to match DataFlow createJob Quota.
+	DATA_FLOW_RL.Take()
+
+	fmt.Println("Launching dataflow job ", dataflowCfg.JobName, " in ", project, "-", dataflowCfg.Location)
 	// Fetch the GCS path from the destination connection profile.
 	dstProf := fmt.Sprintf("projects/%s/locations/%s/connectionProfiles/%s", project, datastreamCfg.DestinationConnectionConfig.Location, datastreamCfg.DestinationConnectionConfig.Name)
-	res, err := dsClient.GetConnectionProfile(ctx, &datastreampb.GetConnectionProfileRequest{Name: dstProf})
+	res, err := datastreamAccessor.GetConnectionProfile(ctx, &datastreampb.GetConnectionProfileRequest{Name: dstProf})
 	if err != nil {
 		return internal.DataflowOutput{}, fmt.Errorf("could not get connection profiles: %v", err)
 	}
@@ -716,8 +740,7 @@
 		Location:        dataflowCfg.Location,
 	}
 	fmt.Println("Created flex template request body...")
-
-	respDf, err := c.LaunchFlexTemplate(ctx, req)
+	respDf, err := dataflowAccessor.LaunchFlexTemplate(ctx, req)
 	if err != nil {
 		fmt.Printf("flexTemplateRequest: %+v\n", req)
 		return internal.DataflowOutput{}, fmt.Errorf("unable to launch template: %v", err)
@@ -877,7 +900,7 @@
 	if err != nil {
 		return internal.DataflowOutput{}, fmt.Errorf("error while writing to GCS: %v", err)
 	}
-	dfOutput, err := LaunchDataflowJob(ctx, targetProfile, streamingCfg, conv)
+	dfOutput, err := LaunchDataflowJob(ctx, targetProfile, streamingCfg, conv, NewDataflowAccessor(), NewDatastreamAccessor())
 	if err != nil {
 		return internal.DataflowOutput{}, fmt.Errorf("error launching dataflow: %v", err)
 	}

streaming/streaming_test.go

@@ -14,15 +14,73 @@
 package streaming
 
 import (
+	"context"
+	"encoding/json"
+	"errors"
+	"io/ioutil"
+	"path/filepath"
 	"sort"
 	"testing"
 
 	"github.com/GoogleCloudPlatform/spanner-migration-tool/internal"
 	"github.com/GoogleCloudPlatform/spanner-migration-tool/profiles"
+	"github.com/googleapis/gax-go/v2"
 	"github.com/stretchr/testify/assert"
 	datastreampb "google.golang.org/genproto/googleapis/cloud/datastream/v1"
+	dataflowpb "google.golang.org/genproto/googleapis/dataflow/v1beta3"
 )
 
+type DataflowAccessorMock struct {}
+var launchFlexTemplateMock func(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error)
+func (dataflowMock DataflowAccessorMock) LaunchFlexTemplate(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error) {
+	return launchFlexTemplateMock(ctx, req)
+}
+
+type DatastreamAccessorMock struct {}
+var getConnectionProfileMock func(ctx context.Context, req *datastreampb.GetConnectionProfileRequest, opts ...gax.CallOption) (*datastreampb.ConnectionProfile, error)
+func (datastreamMocl DatastreamAccessorMock) GetConnectionProfile(ctx context.Context, req *datastreampb.GetConnectionProfileRequest, opts ...gax.CallOption) (*datastreampb.ConnectionProfile, error) {
+	return getConnectionProfileMock(ctx, req)
+}
+
+func TestLaunchDataflowjobWithError(t *testing.T) {
+	Sp := profiles.TargetProfileConnectionSpanner{
+		Endpoint: "zz",
+		Project:  "aa",
+		Instance: "bb",
+		Dbname:   "cc",
+		Dialect:  "dd",
+	}
+	conn := profiles.TargetProfileConnection{Ty: profiles.TargetProfileConnectionTypeSpanner, Sp: Sp}
+	targetProfile := profiles.TargetProfile{
+		Ty:   profiles.TargetProfileTypeConnection,
+		Conn: conn,
+	}
+	path := filepath.Join("..", "test_data", "streamingcfg.json")
+	cfgFile, _ := ioutil.ReadFile(path)
+	streamingCfg := StreamingCfg{}
+	json.Unmarshal(cfgFile, &streamingCfg)
+	dfMock := DataflowAccessorMock{}
+	dsMock := DatastreamAccessorMock{}
+	launchFlexTemplateMock = func(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error) {
+		return &dataflowpb.LaunchFlexTemplateResponse{
+			Job: nil,
+		}, errors.New("this is a mocked error")
+	}
+	getConnectionProfileMock = func(ctx context.Context, req *datastreampb.GetConnectionProfileRequest, opts ...gax.CallOption) (*datastreampb.ConnectionProfile, error) {
+		gcsProfile := datastreampb.GcsProfile{
+			Bucket: "test",
+			RootPath: "test",
+		}
+		return &datastreampb.ConnectionProfile{
+			Name: "test",
+			Profile: &datastreampb.ConnectionProfile_GcsProfile{GcsProfile: &gcsProfile,
+			},
+		}, nil
+	}
+	_, err := LaunchDataflowJob(context.Background(), targetProfile, streamingCfg, nil, dfMock, dsMock)
+	assert.Equal(t, "unable to launch template: this is a mocked error", err.Error())
+}
+
 func TestGetPostgreSQLSourceStreamConfig(t *testing.T) {
 	testCases := []struct {
 		name        string

test_data/streamingcfg.json

@@ -0,0 +1,48 @@
+{
+    "datastreamCfg": {
+     "streamId": "",
+     "streamLocation": "us-central1",
+     "streamDisplayName": "",
+     "sourceConnectionConfig": {
+      "Name": "stalltestxx",
+      "Location": "us-central1"
+     },
+     "destinationConnectionConfig": {
+      "Name": "stalltestyy3",
+      "Location": "us-central1",
+      "Prefix": ""
+     },
+     "properties": "",
+     "maxConcurrentBackfillTasks": "",
+     "maxConcurrentCdcTasks": ""
+    },
+    "gcsCfg": {
+     "ttlInDays": 0,
+     "ttlInDaysSet": false
+    },
+    "dataflowCfg": {
+     "projectId": "",
+     "jobName": "",
+     "location": "us-central1",
+     "hostProjectId": "",
+     "network": "",
+     "subnetwork": "",
+     "maxWorkers": "",
+     "numWorkers": "",
+     "serviceAccountEmail": "",
+     "machineType": "",
+     "additionalUserLabels": "",
+     "kmsKeyName": "",
+     "gcsTemplatePath": "",
+     "dbNameToShardIdMap": null
+    },
+    "tmpDir": "gs://smt-job-6a99-2bcc/",
+    "pubsubCfg": {
+     "TopicId": "",
+     "SubscriptionId": "",
+     "NotificationId": "",
+     "BucketName": "",
+     "Region": ""
+    },
+    "dataShardId": ""
+   }