-
Notifications
You must be signed in to change notification settings - Fork 12
/
state_resolver.go
44 lines (35 loc) · 1.32 KB
/
state_resolver.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
package ledgertests
import (
"github.com/iotaledger/hive.go/ds/shrinkingmap"
"github.com/iotaledger/hive.go/ierrors"
"github.com/iotaledger/iota-core/pkg/core/promise"
"github.com/iotaledger/iota-core/pkg/protocol/engine/mempool"
)
type MockStateResolver struct {
statesByID *shrinkingmap.ShrinkingMap[mempool.StateID, mempool.State]
}
func New(initialStates ...mempool.State) *MockStateResolver {
stateResolver := &MockStateResolver{
statesByID: shrinkingmap.New[mempool.StateID, mempool.State](),
}
for _, initialState := range initialStates {
stateResolver.statesByID.Set(initialState.StateID(), initialState)
}
return stateResolver
}
func (s *MockStateResolver) AddOutputState(state mempool.State) {
s.statesByID.Set(state.StateID(), state)
}
func (s *MockStateResolver) DestroyOutputState(stateID mempool.StateID) {
s.statesByID.Delete(stateID)
}
func (s *MockStateResolver) ResolveOutputState(reference mempool.StateReference) *promise.Promise[mempool.State] {
output, exists := s.statesByID.Get(reference.ReferencedStateID())
if !exists {
return promise.New[mempool.State]().Reject(ierrors.Errorf("output %s not found: %w", reference.ReferencedStateID().ToHex(), mempool.ErrStateNotFound))
}
return promise.New[mempool.State]().Resolve(output)
}
func (s *MockStateResolver) Cleanup() {
s.statesByID.Clear()
}