From da53f35b364d68316a1fb2935fe4eb3257175e92 Mon Sep 17 00:00:00 2001 From: NikitaSkrynnik Date: Tue, 12 Nov 2024 16:07:18 +1100 Subject: [PATCH] Fix review comments Signed-off-by: NikitaSkrynnik --- pkg/networkservice/chains/client/options.go | 4 ++-- pkg/networkservice/chains/nsmgr/heal_test.go | 19 ++++++++----------- .../common/begin/event_factory.go | 12 ++++++------ 3 files changed, 16 insertions(+), 19 deletions(-) diff --git a/pkg/networkservice/chains/client/options.go b/pkg/networkservice/chains/client/options.go index 69a9c3970..d3f0f4564 100644 --- a/pkg/networkservice/chains/client/options.go +++ b/pkg/networkservice/chains/client/options.go @@ -112,8 +112,8 @@ func WithoutRefresh() Option { } } -// WithoutReselectFunc sets a function for changing request parameters on reselect -func WithoutReselectFunc(f func(*networkservice.NetworkServiceRequest)) Option { +// WithReselectFunc sets a function for changing request parameters on reselect +func WithReselectFunc(f func(*networkservice.NetworkServiceRequest)) Option { return func(c *clientOptions) { c.reselectFunc = f } diff --git a/pkg/networkservice/chains/nsmgr/heal_test.go b/pkg/networkservice/chains/nsmgr/heal_test.go index d0731e2c5..0c6f9f8bb 100644 --- a/pkg/networkservice/chains/nsmgr/heal_test.go +++ b/pkg/networkservice/chains/nsmgr/heal_test.go @@ -923,35 +923,32 @@ func TestNSMGRHealEndpoint_CustomReselectFunc(t *testing.T) { SetRegistryProxySupplier(nil). Build() - nsRegistryClient := domain.NewNSRegistryClient(ctx, sandbox.GenerateTestToken) - nsReg, err := nsRegistryClient.Register(ctx, defaultRegistryService(t.Name())) + nsReg, err := domain.NewNSRegistryClient(ctx, sandbox.GenerateTestToken).Register(ctx, defaultRegistryService(t.Name())) require.NoError(t, err) nseReg := defaultRegistryEndpoint(nsReg.Name) nse := domain.Nodes[0].NewEndpoint(ctx, nseReg, sandbox.GenerateTestToken) - request := defaultRequest(nsReg.Name) - nsc := domain.Nodes[0].NewClient(ctx, sandbox.GenerateTestToken, nsclient.WithHealClient(heal.NewClient(ctx)), - nsclient.WithoutReselectFunc( + nsclient.WithReselectFunc( func(request *networkservice.NetworkServiceRequest) { request.Connection.Labels = make(map[string]string) request.Connection.Labels["key"] = "value" request.Connection.NetworkServiceEndpointName = "" })) + request := defaultRequest(nsReg.Name) _, err = nsc.Request(ctx, request.Clone()) require.NoError(t, err) nse.Cancel() - checker := checkrequest.NewServer(t, func(t *testing.T, nsr *networkservice.NetworkServiceRequest) { - require.Contains(t, nsr.Connection.Labels, "key") - }) - counter := new(count.Server) nseReg2 := defaultRegistryEndpoint(nsReg.Name) nseReg2.Name += "-2" - domain.Nodes[0].NewEndpoint(ctx, nseReg2, sandbox.GenerateTestToken, counter, checker) + domain.Nodes[0].NewEndpoint(ctx, nseReg2, sandbox.GenerateTestToken) - require.Eventually(t, func() bool { return counter.UniqueRequests() > 0 }, timeout, tick) + require.Eventually(t, func() bool { + resp, err := nsc.Request(ctx, request.Clone()) + return err == nil && resp.Labels["key"] == "value" + }, timeout, tick) } diff --git a/pkg/networkservice/common/begin/event_factory.go b/pkg/networkservice/common/begin/event_factory.go index 3dd77dbf2..bafa58ec6 100644 --- a/pkg/networkservice/common/begin/event_factory.go +++ b/pkg/networkservice/common/begin/event_factory.go @@ -45,9 +45,11 @@ type ReselectFunc func(request *networkservice.NetworkServiceRequest) // DefaultReselectFunc - default ReselectFunc var DefaultReselectFunc ReselectFunc = func(request *networkservice.NetworkServiceRequest) { - request.GetConnection().Mechanism = nil - request.GetConnection().NetworkServiceEndpointName = "" - request.GetConnection().State = networkservice.State_RESELECT_REQUESTED + if request.GetConnection() != nil { + request.GetConnection().Mechanism = nil + request.GetConnection().NetworkServiceEndpointName = "" + request.GetConnection().State = networkservice.State_RESELECT_REQUESTED + } } // EventFactory - allows firing off a Request or Close event from midchain @@ -115,9 +117,7 @@ func (f *eventFactoryClient) Request(opts ...Option) <-chan error { if o.reselect { ctx, cancel := f.ctxFunc() _, _ = f.client.Close(ctx, request.GetConnection(), f.opts...) - if request.GetConnection() != nil { - f.reselectFunc(request) - } + f.reselectFunc(request) cancel() } ctx, cancel := f.ctxFunc()