Skip to content

Commit

Permalink
[BEAM-13909] improve coverage of Provision package (#17014)
Browse files Browse the repository at this point in the history
* improve coverage of provision package

* updated comments
  • Loading branch information
riteshghorse authored Mar 8, 2022
1 parent 9833b7b commit 0e4d501
Showing 1 changed file with 56 additions and 0 deletions.
56 changes: 56 additions & 0 deletions sdks/go/pkg/beam/provision/provision_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,16 @@
package provision

import (
"context"
"fmt"
"log"
"net"
"reflect"
"sync"
"testing"

fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
"google.golang.org/grpc"
)

type s struct {
Expand Down Expand Up @@ -52,3 +60,51 @@ 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, wg *sync.WaitGroup) {
l, err := net.Listen("tcp", ":0")
defer l.Close()
if err != nil {
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) {

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)
}
want := &fnpb.ProvisionInfo{RetrievalToken: "token"}
if got.GetRetrievalToken() != want.GetRetrievalToken() {
t.Errorf("provision.Info() = %v, want %v", got, want)
}
}

0 comments on commit 0e4d501

Please sign in to comment.