From ba2a7e44179f6239aa9cd1d6db603a0979de366e Mon Sep 17 00:00:00 2001
From: John Guo <john@johng.cn>
Date: Tue, 4 Jul 2023 20:21:12 +0800
Subject: [PATCH] add chaining function Discovery to disable/enable discovery
 feature for package gclient; fix issue #2737 (#2738)

---
 contrib/registry/file/file_z_http_test.go | 37 +++++++++++++++++++++++
 net/gclient/gclient_chain.go              | 10 ++++++
 2 files changed, 47 insertions(+)

diff --git a/contrib/registry/file/file_z_http_test.go b/contrib/registry/file/file_z_http_test.go
index 0401ca8af..57b79c371 100644
--- a/contrib/registry/file/file_z_http_test.go
+++ b/contrib/registry/file/file_z_http_test.go
@@ -12,6 +12,8 @@ import (
 	"time"
 
 	"github.com/gogf/gf/contrib/registry/file/v2"
+	"github.com/gogf/gf/v2/errors/gcode"
+	"github.com/gogf/gf/v2/errors/gerror"
 	"github.com/gogf/gf/v2/frame/g"
 	"github.com/gogf/gf/v2/net/ghttp"
 	"github.com/gogf/gf/v2/net/gsvc"
@@ -48,3 +50,38 @@ func Test_HTTP_Registry(t *testing.T) {
 		t.Assert(client.GetContent(ctx, "/http-registry"), svcName)
 	})
 }
+
+func Test_HTTP_Discovery_Disable(t *testing.T) {
+	var (
+		svcName = guid.S()
+		dirPath = gfile.Temp(svcName)
+	)
+	defer gfile.Remove(dirPath)
+	gsvc.SetRegistry(file.New(dirPath))
+
+	s := g.Server()
+	s.BindHandler("/http-registry", func(r *ghttp.Request) {
+		r.Response.Write(svcName)
+	})
+	s.SetDumpRouterMap(false)
+	s.Start()
+	defer s.Shutdown()
+
+	time.Sleep(100 * time.Millisecond)
+
+	gtest.C(t, func(t *gtest.T) {
+		client := g.Client()
+		client.SetPrefix(fmt.Sprintf("http://127.0.0.1:%d", s.GetListenedPort()))
+		result, err := client.Get(ctx, "/http-registry")
+		defer result.Close()
+		t.Assert(gerror.Code(err), gcode.CodeNotFound)
+	})
+	gtest.C(t, func(t *gtest.T) {
+		client := g.Client()
+		client.SetPrefix(fmt.Sprintf("http://127.0.0.1:%d", s.GetListenedPort()))
+		result, err := client.Discovery(nil).Get(ctx, "/http-registry")
+		defer result.Close()
+		t.AssertNil(err)
+		t.Assert(result.ReadAllString(), svcName)
+	})
+}
diff --git a/net/gclient/gclient_chain.go b/net/gclient/gclient_chain.go
index 03b046936..54926de58 100644
--- a/net/gclient/gclient_chain.go
+++ b/net/gclient/gclient_chain.go
@@ -8,6 +8,8 @@ package gclient
 
 import (
 	"time"
+
+	"github.com/gogf/gf/v2/net/gsvc"
 )
 
 // Prefix is a chaining function,
@@ -37,6 +39,14 @@ func (c *Client) HeaderRaw(headers string) *Client {
 	return newClient
 }
 
+// Discovery is a chaining function, which sets the discovery for client.
+// You can use `Discovery(nil)` to disable discovery feature for current client.
+func (c *Client) Discovery(discovery gsvc.Discovery) *Client {
+	newClient := c.Clone()
+	newClient.SetDiscovery(discovery)
+	return newClient
+}
+
 // Cookie is a chaining function,
 // which sets cookie items with map for next request.
 func (c *Client) Cookie(m map[string]string) *Client {