From 8e48acaeba2758dda9bff9852231a7278af0e03a Mon Sep 17 00:00:00 2001
From: riteshghorse <riteshghorse@gmail.com>
Date: Fri, 4 Mar 2022 05:31:21 +0000
Subject: [PATCH 1/2] improve coverage of provision package

---
 sdks/go/pkg/beam/provision/provision_test.go | 40 ++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/sdks/go/pkg/beam/provision/provision_test.go b/sdks/go/pkg/beam/provision/provision_test.go
index f29bc9b3be51..0b44f256de8b 100644
--- a/sdks/go/pkg/beam/provision/provision_test.go
+++ b/sdks/go/pkg/beam/provision/provision_test.go
@@ -16,8 +16,14 @@
 package provision
 
 import (
+	"context"
+	"log"
+	"net"
 	"reflect"
 	"testing"
+
+	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
+	"google.golang.org/grpc"
 )
 
 type s struct {
@@ -52,3 +58,37 @@ func TestConversions(t *testing.T) {
 		}
 	}
 }
+
+type ProvisionServiceServicer struct {
+	fnpb.UnimplementedProvisionServiceServer
+}
+
+func (p ProvisionServiceServicer) GetProvisionInfo(ctx context.Context, req *fnpb.GetProvisionInfoRequest) (*fnpb.GetProvisionInfoResponse, error) {
+	return &fnpb.GetProvisionInfoResponse{Info: &fnpb.ProvisionInfo{RetrievalToken: "token"}}, nil
+}
+
+func setup(addr string, prs *ProvisionServiceServicer) {
+	l, err := net.Listen("tcp", addr)
+	if err != nil {
+		log.Fatalf("failed to listen on addr: %v", err)
+	}
+	server := grpc.NewServer()
+	defer server.Stop()
+	fnpb.RegisterProvisionServiceServer(server, prs)
+	if err := server.Serve(l); err != nil {
+		log.Fatalf("cannot serve the server: %v", err)
+	}
+}
+
+func TestProvisionInfo(t *testing.T) {
+	prs := &ProvisionServiceServicer{}
+	go setup(":9000", prs)
+	got, err := Info(context.Background(), ":9000")
+	if err != nil {
+		t.Errorf("error in response: %v", err)
+	}
+	want := &fnpb.ProvisionInfo{RetrievalToken: "token"}
+	if got.GetRetrievalToken() != want.GetRetrievalToken() {
+		t.Errorf("provision.Info() = %v, want %v", got, want)
+	}
+}

From 935bc3cc83ab84ab22b037a7ff8b067c5c07c05d Mon Sep 17 00:00:00 2001
From: riteshghorse <riteshghorse@gmail.com>
Date: Fri, 4 Mar 2022 17:23:56 +0000
Subject: [PATCH 2/2] updated comments

---
 sdks/go/pkg/beam/provision/provision_test.go | 28 +++++++++++++++-----
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/sdks/go/pkg/beam/provision/provision_test.go b/sdks/go/pkg/beam/provision/provision_test.go
index 0b44f256de8b..92dc315fc062 100644
--- a/sdks/go/pkg/beam/provision/provision_test.go
+++ b/sdks/go/pkg/beam/provision/provision_test.go
@@ -17,9 +17,11 @@ package provision
 
 import (
 	"context"
+	"fmt"
 	"log"
 	"net"
 	"reflect"
+	"sync"
 	"testing"
 
 	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
@@ -67,23 +69,37 @@ func (p ProvisionServiceServicer) GetProvisionInfo(ctx context.Context, req *fnp
 	return &fnpb.GetProvisionInfoResponse{Info: &fnpb.ProvisionInfo{RetrievalToken: "token"}}, nil
 }
 
-func setup(addr string, prs *ProvisionServiceServicer) {
-	l, err := net.Listen("tcp", addr)
+func setup(addr *string, wg *sync.WaitGroup) {
+	l, err := net.Listen("tcp", ":0")
+	defer l.Close()
 	if err != nil {
-		log.Fatalf("failed to listen on addr: %v", err)
+		log.Fatalf("failed to find an open port: %v", err)
 	}
+	port := l.Addr().(*net.TCPAddr).Port
+	*addr = fmt.Sprintf(":%d", port)
+
 	server := grpc.NewServer()
 	defer server.Stop()
+
+	prs := &ProvisionServiceServicer{}
 	fnpb.RegisterProvisionServiceServer(server, prs)
+
+	wg.Done()
+
 	if err := server.Serve(l); err != nil {
 		log.Fatalf("cannot serve the server: %v", err)
 	}
 }
 
 func TestProvisionInfo(t *testing.T) {
-	prs := &ProvisionServiceServicer{}
-	go setup(":9000", prs)
-	got, err := Info(context.Background(), ":9000")
+
+	endpoint := ""
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go setup(&endpoint, &wg)
+	wg.Wait()
+
+	got, err := Info(context.Background(), endpoint)
 	if err != nil {
 		t.Errorf("error in response: %v", err)
 	}