github.com/1aal/kubeblocks@v0.0.0-20231107070852-e1c03e598921/pkg/lorry/client/suite_test.go (about)

     1  /*
     2  Copyright (C) 2022-2023 ApeCloud Co., Ltd
     3  
     4  This file is part of KubeBlocks project
     5  
     6  This program is free software: you can redistribute it and/or modify
     7  it under the terms of the GNU Affero General Public License as published by
     8  the Free Software Foundation, either version 3 of the License, or
     9  (at your option) any later version.
    10  
    11  This program is distributed in the hope that it will be useful
    12  but WITHOUT ANY WARRANTY; without even the implied warranty of
    13  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    14  GNU Affero General Public License for more details.
    15  
    16  You should have received a copy of the GNU Affero General Public License
    17  along with this program.  If not, see <http://www.gnu.org/licenses/>.
    18  */
    19  
    20  package client
    21  
    22  import (
    23  	"fmt"
    24  	"net"
    25  	"testing"
    26  
    27  	"github.com/golang/mock/gomock"
    28  	. "github.com/onsi/ginkgo/v2"
    29  	. "github.com/onsi/gomega"
    30  	"github.com/spf13/viper"
    31  	"github.com/valyala/fasthttp"
    32  
    33  	"github.com/1aal/kubeblocks/pkg/lorry/dcs"
    34  	"github.com/1aal/kubeblocks/pkg/lorry/engines"
    35  	"github.com/1aal/kubeblocks/pkg/lorry/engines/register"
    36  	"github.com/1aal/kubeblocks/pkg/lorry/httpserver"
    37  	opsregister "github.com/1aal/kubeblocks/pkg/lorry/operations/register"
    38  )
    39  
    40  var (
    41  	lorryHTTPPort = 3501
    42  	tcpListener   net.Listener
    43  	dbManager     engines.DBManager
    44  	mockDBManager *engines.MockDBManager
    45  	dcsStore      dcs.DCS
    46  	mockDCSStore  *dcs.MockDCS
    47  )
    48  
    49  func init() {
    50  	viper.AutomaticEnv()
    51  	viper.SetDefault("KB_POD_NAME", "pod-test")
    52  	viper.SetDefault("KB_CLUSTER_COMP_NAME", "cluster-component-test")
    53  	viper.SetDefault("KB_NAMESPACE", "namespace-test")
    54  }
    55  
    56  func TestLorryClient(t *testing.T) {
    57  	RegisterFailHandler(Fail)
    58  	RunSpecs(t, "Lorry Client. Suite")
    59  }
    60  
    61  var _ = BeforeSuite(func() {
    62  	// Init mock db manager
    63  	InitMockDBManager()
    64  
    65  	// Init mock dcs store
    66  	InitMockDCSStore()
    67  
    68  	// Start lorry HTTP server
    69  	StartHTTPServer()
    70  })
    71  
    72  var _ = AfterSuite(func() {
    73  	StopHTTPServer()
    74  })
    75  
    76  func newTCPServer(port int) (net.Listener, int) {
    77  	for i := 0; i < 3; i++ {
    78  		tcpListener, _ = net.Listen("tcp", fmt.Sprintf(":%v", port))
    79  		if tcpListener != nil {
    80  			break
    81  		}
    82  		port++
    83  	}
    84  	Expect(tcpListener).ShouldNot(BeNil())
    85  	return tcpListener, port
    86  }
    87  
    88  func StartHTTPServer() {
    89  	ops := opsregister.Operations()
    90  	s := httpserver.NewServer(ops)
    91  	handler := s.Router()
    92  
    93  	listener, port := newTCPServer(lorryHTTPPort)
    94  	lorryHTTPPort = port
    95  
    96  	go func() {
    97  		Expect(fasthttp.Serve(listener, handler)).Should(Succeed())
    98  	}()
    99  }
   100  
   101  func StopHTTPServer() {
   102  	if tcpListener == nil {
   103  		return
   104  	}
   105  	_ = tcpListener.Close()
   106  }
   107  
   108  func InitMockDBManager() {
   109  	ctrl := gomock.NewController(GinkgoT())
   110  	mockDBManager = engines.NewMockDBManager(ctrl)
   111  	register.SetDBManager(mockDBManager)
   112  	dbManager = mockDBManager
   113  }
   114  
   115  func InitMockDCSStore() {
   116  	ctrl := gomock.NewController(GinkgoT())
   117  	mockDCSStore = dcs.NewMockDCS(ctrl)
   118  	mockDCSStore.EXPECT().GetClusterFromCache().Return(&dcs.Cluster{}).AnyTimes()
   119  	dcs.SetStore(mockDCSStore)
   120  	dcsStore = mockDCSStore
   121  }