github.com/MetalBlockchain/metalgo@v1.11.9/snow/uptime/test_state.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package uptime
     5  
     6  import (
     7  	"time"
     8  
     9  	"github.com/MetalBlockchain/metalgo/database"
    10  	"github.com/MetalBlockchain/metalgo/ids"
    11  )
    12  
    13  var _ State = (*TestState)(nil)
    14  
    15  type uptime struct {
    16  	upDuration  time.Duration
    17  	lastUpdated time.Time
    18  	startTime   time.Time
    19  }
    20  
    21  type TestState struct {
    22  	dbReadError  error
    23  	dbWriteError error
    24  	nodes        map[ids.NodeID]map[ids.ID]*uptime
    25  }
    26  
    27  func NewTestState() *TestState {
    28  	return &TestState{
    29  		nodes: make(map[ids.NodeID]map[ids.ID]*uptime),
    30  	}
    31  }
    32  
    33  func (s *TestState) AddNode(nodeID ids.NodeID, subnetID ids.ID, startTime time.Time) {
    34  	subnetUptimes, ok := s.nodes[nodeID]
    35  	if !ok {
    36  		subnetUptimes = make(map[ids.ID]*uptime)
    37  		s.nodes[nodeID] = subnetUptimes
    38  	}
    39  	st := time.Unix(startTime.Unix(), 0)
    40  	subnetUptimes[subnetID] = &uptime{
    41  		lastUpdated: st,
    42  		startTime:   st,
    43  	}
    44  }
    45  
    46  func (s *TestState) GetUptime(nodeID ids.NodeID, subnetID ids.ID) (time.Duration, time.Time, error) {
    47  	up, exists := s.nodes[nodeID][subnetID]
    48  	if !exists {
    49  		return 0, time.Time{}, database.ErrNotFound
    50  	}
    51  	return up.upDuration, up.lastUpdated, s.dbReadError
    52  }
    53  
    54  func (s *TestState) SetUptime(nodeID ids.NodeID, subnetID ids.ID, upDuration time.Duration, lastUpdated time.Time) error {
    55  	up, exists := s.nodes[nodeID][subnetID]
    56  	if !exists {
    57  		return database.ErrNotFound
    58  	}
    59  	up.upDuration = upDuration
    60  	up.lastUpdated = time.Unix(lastUpdated.Unix(), 0)
    61  	return s.dbWriteError
    62  }
    63  
    64  func (s *TestState) GetStartTime(nodeID ids.NodeID, subnetID ids.ID) (time.Time, error) {
    65  	up, exists := s.nodes[nodeID][subnetID]
    66  	if !exists {
    67  		return time.Time{}, database.ErrNotFound
    68  	}
    69  	return up.startTime, s.dbReadError
    70  }