github.com/xmlking/toolkit/broker/pubsub@v0.3.4/broker_test.go (about)

     1  package broker_test
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"cloud.google.com/go/pubsub"
     8  	"cloud.google.com/go/pubsub/pstest"
     9  	broker "github.com/xmlking/toolkit/broker/pubsub"
    10  	"google.golang.org/api/option"
    11  	pb "google.golang.org/genproto/googleapis/pubsub/v1"
    12  	"google.golang.org/grpc"
    13  )
    14  
    15  func TestNewBroker(t *testing.T) {
    16  	ctX, cancel := context.WithCancel(context.Background())
    17  	srv := setupFakePubsubAndBroker(ctX, t)
    18  	defer srv.Close()
    19  
    20  	myHandler := func(ctx context.Context, msg *pubsub.Message) {
    21  		t.Logf("received msg: %s", msg.Data)
    22  		msg.Ack()
    23  		cancel()
    24  		if got, want := string(msg.Data), "hello"; got != want {
    25  			t.Fatalf(`got "%s" message, want "%s"`, got, want)
    26  			msg.Nack()
    27  		}
    28  	}
    29  
    30  	// add subscriber
    31  	if err := broker.AddSubscriber("sumo", myHandler); err != nil {
    32  		t.Fatal(err)
    33  	}
    34  
    35  	// start broker
    36  	go broker.Start()
    37  
    38  	// create publisher
    39  	publisher, err := broker.NewPublisher("sumo", broker.PublishAsync(false))
    40  	if err != nil {
    41  		t.Fatal(err)
    42  	}
    43  
    44  	// publish message
    45  	err = publisher.Publish(context.Background(), &pubsub.Message{Data: []byte("hello")})
    46  	if err != nil {
    47  		t.Fatal(err)
    48  	}
    49  
    50  	//srv.Wait()
    51  	ms := srv.Messages()
    52  	t.Logf("msg1: %s", ms[0].Data)
    53  	tps, err := srv.GServer.ListTopics(context.TODO(), &pb.ListTopicsRequest{Project: "projects/pro1"})
    54  	if err != nil {
    55  		t.Fatal(err)
    56  	}
    57  	t.Log(tps)
    58  }
    59  
    60  func TestSubscribeWithRecoveryHandler(t *testing.T) {
    61  	ctX, cancel := context.WithCancel(context.Background())
    62  	srv := setupFakePubsubAndBroker(ctX, t)
    63  	defer srv.Close()
    64  
    65  	myHandler := func(ctx context.Context, msg *pubsub.Message) {
    66  		t.Logf("received msg: %s", msg.Data)
    67  		msg.Ack()
    68  		cancel()
    69  		if got, want := string(msg.Data), "hello"; got != want {
    70  			t.Fatalf(`got "%s" message, want "%s"`, got, want)
    71  			msg.Nack()
    72  		}
    73  		panic("sumo")
    74  	}
    75  
    76  	recoveryHandler := func(ctx context.Context, msg *pubsub.Message, r interface{}) {
    77  		t.Logf("Recovered from panic: %v,  msg.ID: %s", r, msg.ID)
    78  	}
    79  
    80  	// add subscriber
    81  	if err := broker.AddSubscriber("sumo", myHandler, broker.WithRecoveryHandler(recoveryHandler)); err != nil {
    82  		t.Fatal(err)
    83  	}
    84  
    85  	// start broker
    86  	go broker.Start()
    87  
    88  	// create publisher
    89  	publisher, err := broker.NewPublisher("sumo", broker.PublishAsync(false))
    90  	if err != nil {
    91  		t.Fatal(err)
    92  	}
    93  
    94  	// publish message
    95  	err = publisher.Publish(context.Background(), &pubsub.Message{Data: []byte("hello")})
    96  	if err != nil {
    97  		t.Fatal(err)
    98  	}
    99  
   100  	//srv.Wait()
   101  	ms := srv.Messages()
   102  	t.Logf("msg1: %s", ms[0].Data)
   103  	tps, err := srv.GServer.ListTopics(context.TODO(), &pb.ListTopicsRequest{Project: "projects/pro1"})
   104  	if err != nil {
   105  		t.Fatal(err)
   106  	}
   107  	t.Log(tps)
   108  }
   109  
   110  // setupFakePubsubAndBroker creates a new fake pubsub server and setup topics and subscriptions
   111  // it also created broker and set as default.
   112  // Note: be sure to close server and broker!
   113  func setupFakePubsubAndBroker(ctx context.Context, t *testing.T) *pstest.Server {
   114  	t.Helper()
   115  
   116  	srv := pstest.NewServer()
   117  	// Connect to the server without using TLS.
   118  	conn, err := grpc.DialContext(ctx, srv.Addr, grpc.WithInsecure())
   119  	if err != nil {
   120  		t.Fatal(err)
   121  	}
   122  
   123  	client, err := pubsub.NewClient(ctx, "pro1", option.WithGRPCConn(conn))
   124  	if err != nil {
   125  		t.Fatal(err)
   126  	}
   127  
   128  	ctxl := context.Background()
   129  
   130  	// srv.GServer.CreateTopic(ctxl, &pb.Topic{Name:"projects/pro1/topics/sumo"})
   131  	// create topic "projects/pro1/topics/sumo"
   132  	top, err := client.CreateTopic(ctxl, "sumo")
   133  	if err != nil {
   134  		t.Fatal(err)
   135  	}
   136  
   137  	// srv.GServer.CreateSubscription(ctxl, &pb.Subscription{Name:"projects/pro1/subscriptions/sumo", Topic: top.String()})
   138  	// create subscription "projects/pro1/subscriptions/sumo"
   139  	sub, err := client.CreateSubscription(ctxl, "sumo", pubsub.SubscriptionConfig{Topic: top})
   140  	if err != nil {
   141  		t.Fatal(err)
   142  	}
   143  	t.Log(sub)
   144  
   145  	broker.DefaultBroker = broker.NewBroker(ctx, broker.ProjectID("pro1"), broker.ClientOption(option.WithGRPCConn(conn)))
   146  	return srv
   147  }