github.com/kubeshop/testkube@v1.17.23/pkg/logs/state/state.go (about)

     1  package state
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  
     7  	"github.com/nats-io/nats.go/jetstream"
     8  )
     9  
    10  var (
    11  	// state not found error
    12  	ErrStateNotFound = errors.New("no state found")
    13  )
    14  
    15  // NewState creates new state storage
    16  func NewState(kv jetstream.KeyValue) Interface {
    17  	return &State{
    18  		kv: kv,
    19  	}
    20  }
    21  
    22  // State is a state storage based on NATS KV store
    23  type State struct {
    24  	kv jetstream.KeyValue
    25  }
    26  
    27  // Get returns state for given key - executionId
    28  func (s State) Get(ctx context.Context, key string) (LogState, error) {
    29  	state, err := s.kv.Get(ctx, key)
    30  	if err != nil {
    31  		if err == jetstream.ErrKeyNotFound {
    32  			return LogStateUnknown, nil
    33  		}
    34  
    35  		return LogStateUnknown, err
    36  	}
    37  
    38  	if len(state.Value()) == 0 {
    39  		return LogStateUnknown, ErrStateNotFound
    40  	}
    41  
    42  	return LogState(state.Value()[0]), nil
    43  }
    44  
    45  // Put puts state for given key - executionId
    46  func (s State) Put(ctx context.Context, key string, state LogState) error {
    47  	_, err := s.kv.Put(ctx, key, []byte{byte(state)})
    48  	return err
    49  }