From 7be26744e042e745c1afb4ef5566f2a6c8dc2908 Mon Sep 17 00:00:00 2001
From: fujiwara <fujiwara.shunichiro@gmail.com>
Date: Tue, 24 Oct 2023 17:55:47 +0900
Subject: [PATCH] add gRPC client

---
 client.go    |  5 +++++
 grpc.go      | 35 +++++++++++++++++++++++++++++++++++
 grpc_test.go | 43 +++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/client.go b/client.go
index 954015a..632574d 100644
--- a/client.go
+++ b/client.go
@@ -12,6 +12,11 @@ import (
 // DefaultClientTimeout is default timeout for katsubushi client
 var DefaultClientTimeout = 5 * time.Second
 
+type ClientInterface interface {
+	Fetch(ctx context.Context) (uint64, error)
+	FetchMulti(ctx context.Context, n int) ([]uint64, error)
+}
+
 // Client is katsubushi client
 type Client struct {
 	memcacheClients []*memcacheClient
diff --git a/grpc.go b/grpc.go
index 85e365b..79259b7 100644
--- a/grpc.go
+++ b/grpc.go
@@ -13,6 +13,7 @@ import (
 	grpc_recovery "github.com/grpc-ecosystem/go-grpc-middleware/recovery"
 	gogrpc "google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/reflection"
 	"google.golang.org/grpc/status"
 )
@@ -119,3 +120,37 @@ func (sv *gRPCStats) Get(ctx context.Context, req *grpc.StatsRequest) (*grpc.Sta
 		GetMisses:        st.GetMisses,
 	}, nil
 }
+
+type GRPCClient struct {
+	client grpc.GeneratorClient
+}
+
+func NewGRPCClient(addr string) (*GRPCClient, error) {
+	conn, err := gogrpc.Dial(
+		addr,
+		gogrpc.WithTransportCredentials(insecure.NewCredentials()),
+		gogrpc.WithBlock(),
+	)
+	if err != nil {
+		return nil, err
+	}
+	return &GRPCClient{
+		client: grpc.NewGeneratorClient(conn),
+	}, nil
+}
+
+func (g *GRPCClient) Fetch(ctx context.Context) (uint64, error) {
+	res, err := g.client.Fetch(ctx, &grpc.FetchRequest{})
+	if err != nil {
+		return 0, err
+	}
+	return res.Id, nil
+}
+
+func (g *GRPCClient) FetchMulti(ctx context.Context, n int) ([]uint64, error) {
+	res, err := g.client.FetchMulti(ctx, &grpc.FetchMultiRequest{N: uint32(n)})
+	if err != nil {
+		return nil, err
+	}
+	return res.Ids, nil
+}
diff --git a/grpc_test.go b/grpc_test.go
index aad60cb..a6cd38a 100644
--- a/grpc_test.go
+++ b/grpc_test.go
@@ -60,7 +60,7 @@ func TestGRPCSingle(t *testing.T) {
 		if res.Id == 0 {
 			t.Fatal("id should not be 0")
 		}
-		t.Logf("HTTP fetched single ID: %d", res.Id)
+		t.Logf("gRPC fetched single ID: %d", res.Id)
 	}
 }
 
@@ -83,7 +83,46 @@ func TestGRPCMulti(t *testing.T) {
 				t.Fatal("id should not be 0")
 			}
 		}
-		t.Logf("HTTP fetched IDs: %v", res.Ids)
+		t.Logf("gRPC fetched IDs: %v", res.Ids)
+	}
+}
+
+func TestGRPCClientSingle(t *testing.T) {
+	client, err := katsubushi.NewGRPCClient(fmt.Sprintf("localhost:%d", grpcPort))
+	if err != nil {
+		t.Fatal(err)
+	}
+	for i := 0; i < 10; i++ {
+		id, err := client.Fetch(context.Background())
+		if err != nil {
+			t.Fatal(err)
+		}
+		if id == 0 {
+			t.Fatal("id should not be 0")
+		}
+		t.Logf("gRPC fetched single ID: %d", id)
+	}
+}
+
+func TestGRPCClientMulti(t *testing.T) {
+	client, err := katsubushi.NewGRPCClient(fmt.Sprintf("localhost:%d", grpcPort))
+	if err != nil {
+		t.Fatal(err)
+	}
+	for i := 0; i < 10; i++ {
+		ids, err := client.FetchMulti(context.Background(), 10)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(ids) != 10 {
+			t.Fatalf("ids should contain 10 elements %v", ids)
+		}
+		for _, id := range ids {
+			if id == 0 {
+				t.Fatal("id should not be 0")
+			}
+		}
+		t.Logf("gRPC fetched IDs: %v", ids)
 	}
 }