github.com/apache/beam/sdks/v2@v2.48.2/go/container/tools/provision_test.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one or more
     2  // contributor license agreements.  See the NOTICE file distributed with
     3  // this work for additional information regarding copyright ownership.
     4  // The ASF licenses this file to You under the Apache License, Version 2.0
     5  // (the "License"); you may not use this file except in compliance with
     6  // the License.  You may obtain a copy of the License at
     7  //
     8  //    http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  package tools
    17  
    18  import (
    19  	"context"
    20  	"fmt"
    21  	"log"
    22  	"net"
    23  	"reflect"
    24  	"sync"
    25  	"testing"
    26  
    27  	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
    28  	"google.golang.org/grpc"
    29  )
    30  
    31  type s struct {
    32  	A int    `json:"a,omitempty"`
    33  	B string `json:"b,omitempty"`
    34  	C bool   `json:"c,omitempty"`
    35  	D *s     `json:"d,omitempty"`
    36  }
    37  
    38  // TestConversions verifies that we can process proto structs via JSON.
    39  func TestConversions(t *testing.T) {
    40  	tests := []s{
    41  		s{},
    42  		s{A: 2},
    43  		s{B: "foo"},
    44  		s{C: true},
    45  		s{D: &s{A: 3}},
    46  		s{A: 1, B: "bar", C: true, D: &s{A: 3, B: "baz"}},
    47  	}
    48  
    49  	for _, test := range tests {
    50  		enc, err := OptionsToProto(test)
    51  		if err != nil {
    52  			t.Errorf("Failed to marshal %v: %v", test, err)
    53  		}
    54  		var ret s
    55  		if err := ProtoToOptions(enc, &ret); err != nil {
    56  			t.Errorf("Failed to unmarshal %v from %v: %v", test, enc, err)
    57  		}
    58  		if !reflect.DeepEqual(test, ret) {
    59  			t.Errorf("Unmarshal(Marshal(%v)) = %v, want %v", test, ret, test)
    60  		}
    61  	}
    62  }
    63  
    64  type ProvisionServiceServicer struct {
    65  	fnpb.UnimplementedProvisionServiceServer
    66  }
    67  
    68  func (p ProvisionServiceServicer) GetProvisionInfo(ctx context.Context, req *fnpb.GetProvisionInfoRequest) (*fnpb.GetProvisionInfoResponse, error) {
    69  	return &fnpb.GetProvisionInfoResponse{Info: &fnpb.ProvisionInfo{RetrievalToken: "token"}}, nil
    70  }
    71  
    72  func setup(addr *string, wg *sync.WaitGroup) {
    73  	l, err := net.Listen("tcp", ":0")
    74  	if err != nil {
    75  		log.Fatalf("failed to find an open port: %v", err)
    76  	}
    77  	defer l.Close()
    78  	port := l.Addr().(*net.TCPAddr).Port
    79  	*addr = fmt.Sprintf(":%d", port)
    80  
    81  	server := grpc.NewServer()
    82  	defer server.Stop()
    83  
    84  	prs := &ProvisionServiceServicer{}
    85  	fnpb.RegisterProvisionServiceServer(server, prs)
    86  
    87  	wg.Done()
    88  
    89  	if err := server.Serve(l); err != nil {
    90  		log.Fatalf("cannot serve the server: %v", err)
    91  	}
    92  }
    93  
    94  func TestProvisionInfo(t *testing.T) {
    95  
    96  	endpoint := ""
    97  	var wg sync.WaitGroup
    98  	wg.Add(1)
    99  	go setup(&endpoint, &wg)
   100  	wg.Wait()
   101  
   102  	got, err := ProvisionInfo(context.Background(), endpoint)
   103  	if err != nil {
   104  		t.Errorf("error in response: %v", err)
   105  	}
   106  	want := &fnpb.ProvisionInfo{RetrievalToken: "token"}
   107  	if got.GetRetrievalToken() != want.GetRetrievalToken() {
   108  		t.Errorf("provision.Info() = %v, want %v", got, want)
   109  	}
   110  }