From 0886d800838fc9877446cffceefeb42cfd31c45a Mon Sep 17 00:00:00 2001
From: Marina Biryukova <m.biryukova@yadro.com>
Date: Fri, 28 Jul 2023 18:35:35 +0300
Subject: [PATCH] [#69] Add close for nns.Dial

Signed-off-by: Marina Biryukova <m.biryukova@yadro.com>
---
 ns/nns.go | 39 +++++++++++++++++++++++----------------
 1 file changed, 23 insertions(+), 16 deletions(-)

diff --git a/ns/nns.go b/ns/nns.go
index 4af72f4..947e8fa 100644
--- a/ns/nns.go
+++ b/ns/nns.go
@@ -19,12 +19,25 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
 )
 
+// multiSchemeClient unites invoker.RPCInvoke and common interface of
+// rpcclient.Client and rpcclient.WSClient.
+type multiSchemeClient interface {
+	invoker.RPCInvoke
+	// Init turns client to "ready-to-work" state.
+	Init() error
+	// Close closes connections.
+	Close()
+	// GetContractStateByID returns state of the NNS contract on 1 input.
+	GetContractStateByID(int32) (*state.Contract, error)
+}
+
 // NNS looks up FrostFS names using Neo Name Service.
 //
 // Instances are created with a variable declaration. Before work, the connection
 // to the NNS server MUST be established using Dial method.
 type NNS struct {
 	nnsContract util.Uint160
+	client      multiSchemeClient
 
 	invoker interface {
 		Call(contract util.Uint160, operation string, params ...any) (*result.Invoke, error)
@@ -37,47 +50,41 @@ type NNS struct {
 // If URL address scheme is 'ws' or 'wss', then WebSocket protocol is used,
 // otherwise HTTP.
 func (n *NNS) Dial(address string) error {
-	// multiSchemeClient unites invoker.RPCInvoke and common interface of
-	// rpcclient.Client and rpcclient.WSClient. Interface is anonymous
-	// according to assumption that common interface of these client types
-	// is not required by design and may diverge with changes.
-	var multiSchemeClient interface {
-		invoker.RPCInvoke
-		// Init turns client to "ready-to-work" state.
-		Init() error
-		// GetContractStateByID returns state of the NNS contract on 1 input.
-		GetContractStateByID(int32) (*state.Contract, error)
-	}
 	var err error
 
 	uri, err := url.Parse(address)
 	if err == nil && (uri.Scheme == "ws" || uri.Scheme == "wss") {
-		multiSchemeClient, err = rpcclient.NewWS(context.Background(), address, rpcclient.WSOptions{})
+		n.client, err = rpcclient.NewWS(context.Background(), address, rpcclient.WSOptions{})
 		if err != nil {
 			return fmt.Errorf("create Neo WebSocket client: %w", err)
 		}
 	} else {
-		multiSchemeClient, err = rpcclient.New(context.Background(), address, rpcclient.Options{})
+		n.client, err = rpcclient.New(context.Background(), address, rpcclient.Options{})
 		if err != nil {
 			return fmt.Errorf("create Neo HTTP client: %w", err)
 		}
 	}
 
-	if err = multiSchemeClient.Init(); err != nil {
+	if err = n.client.Init(); err != nil {
 		return fmt.Errorf("initialize Neo client: %w", err)
 	}
 
-	nnsContract, err := multiSchemeClient.GetContractStateByID(1)
+	nnsContract, err := n.client.GetContractStateByID(1)
 	if err != nil {
 		return fmt.Errorf("get NNS contract state: %w", err)
 	}
 
-	n.invoker = invoker.New(multiSchemeClient, nil)
+	n.invoker = invoker.New(n.client, nil)
 	n.nnsContract = nnsContract.Hash
 
 	return nil
 }
 
+// Close closes connections of multiSchemeClient.
+func (n *NNS) Close() {
+	n.client.Close()
+}
+
 // ResolveContainerDomain looks up for NNS TXT records for the given container domain
 // by calling `resolve` method of NNS contract. Returns the first record which represents
 // valid container ID in a string format. Otherwise, returns an error.