github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/module/irrecoverable/unittest.go (about)

     1  package irrecoverable
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/stretchr/testify/mock"
     8  	"github.com/stretchr/testify/require"
     9  )
    10  
    11  // MockSignalerContext is a SignalerContext that can be used in tests to assert that an error is thrown.
    12  // It embeds a mock.Mock, so it can be used it to assert that Throw is called with a specific error.
    13  // Use NewMockSignalerContextExpectError to create a new MockSignalerContext that expects a specific error, otherwise NewMockSignalerContext.
    14  type MockSignalerContext struct {
    15  	context.Context
    16  	*mock.Mock
    17  }
    18  
    19  var _ SignalerContext = &MockSignalerContext{}
    20  
    21  func (m MockSignalerContext) sealed() {}
    22  
    23  func (m MockSignalerContext) Throw(err error) {
    24  	m.Called(err)
    25  }
    26  
    27  func NewMockSignalerContext(t *testing.T, ctx context.Context) *MockSignalerContext {
    28  	m := &MockSignalerContext{
    29  		Context: ctx,
    30  		Mock:    &mock.Mock{},
    31  	}
    32  	m.Mock.Test(t)
    33  	t.Cleanup(func() { m.AssertExpectations(t) })
    34  	return m
    35  }
    36  
    37  // NewMockSignalerContextWithCancel creates a new MockSignalerContext with a cancel function.
    38  func NewMockSignalerContextWithCancel(t *testing.T, parent context.Context) (*MockSignalerContext, context.CancelFunc) {
    39  	ctx, cancel := context.WithCancel(parent)
    40  	return NewMockSignalerContext(t, ctx), cancel
    41  }
    42  
    43  // NewMockSignalerContextExpectError creates a new MockSignalerContext which expects a specific error to be thrown.
    44  func NewMockSignalerContextExpectError(t *testing.T, ctx context.Context, err error) *MockSignalerContext {
    45  	require.NotNil(t, err)
    46  	m := NewMockSignalerContext(t, ctx)
    47  
    48  	// since we expect an error, we should expect a call to Throw
    49  	m.On("Throw", err).Once().Return()
    50  
    51  	return m
    52  }