github.com/Ilhicas/nomad@v1.0.4-0.20210304152020-e86851182bc3/client/testing.go (about)

     1  package client
     2  
     3  import (
     4  	"fmt"
     5  	"time"
     6  
     7  	"github.com/hashicorp/nomad/client/config"
     8  	consulapi "github.com/hashicorp/nomad/client/consul"
     9  	"github.com/hashicorp/nomad/client/fingerprint"
    10  	agentconsul "github.com/hashicorp/nomad/command/agent/consul"
    11  	"github.com/hashicorp/nomad/helper/pluginutils/catalog"
    12  	"github.com/hashicorp/nomad/helper/pluginutils/singleton"
    13  	"github.com/hashicorp/nomad/helper/testlog"
    14  	testing "github.com/mitchellh/go-testing-interface"
    15  )
    16  
    17  // TestClient creates an in-memory client for testing purposes and returns a
    18  // cleanup func to shutdown the client and remove the alloc and state dirs.
    19  //
    20  // There is no need to override the AllocDir or StateDir as they are randomized
    21  // and removed in the returned cleanup function. If they are overridden in the
    22  // callback then the caller still must run the returned cleanup func.
    23  func TestClient(t testing.T, cb func(c *config.Config)) (*Client, func() error) {
    24  	conf, cleanup := config.TestClientConfig(t)
    25  
    26  	// Tighten the fingerprinter timeouts (must be done in client package
    27  	// to avoid circular dependencies)
    28  	if conf.Options == nil {
    29  		conf.Options = make(map[string]string)
    30  	}
    31  	conf.Options[fingerprint.TightenNetworkTimeoutsConfig] = "true"
    32  
    33  	logger := testlog.HCLogger(t)
    34  	conf.Logger = logger
    35  
    36  	if cb != nil {
    37  		cb(conf)
    38  	}
    39  
    40  	// Set the plugin loaders
    41  	if conf.PluginLoader == nil {
    42  		conf.PluginLoader = catalog.TestPluginLoaderWithOptions(t, "", conf.Options, nil)
    43  	}
    44  	if conf.PluginSingletonLoader == nil {
    45  		conf.PluginSingletonLoader = singleton.NewSingletonLoader(logger, conf.PluginLoader)
    46  	}
    47  	mockCatalog := agentconsul.NewMockCatalog(logger)
    48  	mockService := consulapi.NewMockConsulServiceClient(t, logger)
    49  	client, err := NewClient(conf, mockCatalog, nil, mockService)
    50  	if err != nil {
    51  		cleanup()
    52  		t.Fatalf("err: %v", err)
    53  	}
    54  	return client, func() error {
    55  		ch := make(chan error)
    56  
    57  		go func() {
    58  			defer close(ch)
    59  
    60  			// Shutdown client
    61  			err := client.Shutdown()
    62  			if err != nil {
    63  				ch <- fmt.Errorf("failed to shutdown client: %v", err)
    64  			}
    65  
    66  			// Call TestClientConfig cleanup
    67  			cleanup()
    68  		}()
    69  
    70  		select {
    71  		case e := <-ch:
    72  			return e
    73  		case <-time.After(1 * time.Minute):
    74  			return fmt.Errorf("timed out while shutting down client")
    75  		}
    76  	}
    77  }