github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/scheduler/benchmarks/helpers_test.go (about)

     1  package benchmarks
     2  
     3  // Test helper functions for running scheduling tests and benchmarks
     4  // against real world state snapshots or data directories. These live
     5  // here and not in the the parent scheduler package because it would
     6  // create circular imports between the scheduler and raftutils package
     7  // (via the nomad package)
     8  
     9  import (
    10  	"errors"
    11  	"os"
    12  	"testing"
    13  
    14  	"github.com/hashicorp/nomad/helper/raftutil"
    15  	"github.com/hashicorp/nomad/scheduler"
    16  	"github.com/stretchr/testify/require"
    17  )
    18  
    19  // NewBenchmarkingHarness creates a starting test harness with state
    20  // store. The starting contents of the state store depends on which
    21  // env var is set:
    22  // - NOMAD_BENCHMARK_DATADIR: path to data directory
    23  // - NOMAD_BENCHMARK_SNAPSHOT: path to raft snapshot
    24  // - neither: empty starting state
    25  func NewBenchmarkingHarness(t testing.TB) *scheduler.Harness {
    26  	// create the Harness and starting state.
    27  	datadir := os.Getenv("NOMAD_BENCHMARK_DATADIR")
    28  	if datadir != "" {
    29  		h, err := NewHarnessFromDataDir(t, datadir)
    30  		require.NoError(t, err)
    31  		return h
    32  	} else {
    33  		snapshotPath := os.Getenv("NOMAD_BENCHMARK_SNAPSHOT")
    34  		if snapshotPath != "" {
    35  			h, err := NewHarnessFromSnapshot(t, snapshotPath)
    36  			require.NoError(t, err)
    37  			return h
    38  		}
    39  	}
    40  	return scheduler.NewHarness(t)
    41  }
    42  
    43  // NewHarnessFromDataDir creates a new scheduler test harness with
    44  // state loaded from an existing datadir.
    45  func NewHarnessFromDataDir(t testing.TB, datadirPath string) (*scheduler.Harness, error) {
    46  	if datadirPath == "" {
    47  		return nil, errors.New("datadir path was not set")
    48  	}
    49  	fsm, err := raftutil.NewFSM(datadirPath)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  	_, _, err = fsm.ApplyAll()
    54  	if err != nil {
    55  		return nil, err
    56  	}
    57  
    58  	return scheduler.NewHarnessWithState(t, fsm.State()), nil
    59  }
    60  
    61  // NewHarnessFromDataDir creates a new harness with state loaded
    62  // from an existing raft snapshot.
    63  func NewHarnessFromSnapshot(t testing.TB, snapshotPath string) (*scheduler.Harness, error) {
    64  	if snapshotPath == "" {
    65  		return nil, errors.New("snapshot path was not set")
    66  	}
    67  	f, err := os.Open(snapshotPath)
    68  	if err != nil {
    69  		return nil, err
    70  	}
    71  	defer f.Close()
    72  
    73  	state, _, err := raftutil.RestoreFromArchive(f, nil)
    74  	if err != nil {
    75  		return nil, err
    76  	}
    77  
    78  	return scheduler.NewHarnessWithState(t, state), nil
    79  }