diff --git a/api/v1/validatorstate.go b/api/v1/validatorstate.go index 404b4a9d..f9db17f9 100644 --- a/api/v1/validatorstate.go +++ b/api/v1/validatorstate.go @@ -15,6 +15,7 @@ package v1 import ( "fmt" + "github.com/pkg/errors" "strings" "github.com/attestantio/go-eth2-client/spec/phase0" @@ -94,8 +95,12 @@ func (v *ValidatorState) UnmarshalJSON(input []byte) error { return err } -func (v ValidatorState) String() string { - return validatorStateStrings[v] +func (v ValidatorState) String() (string, error) { + if int(v) < 0 || int(v) > len(validatorStateStrings) { + return "", errors.New("invalid validator state") + } + + return validatorStateStrings[v], nil } // IsPending returns true if the validator is pending. diff --git a/api/v1/validatorstate_test.go b/api/v1/validatorstate_test.go index 00e71efb..7bc6b385 100644 --- a/api/v1/validatorstate_test.go +++ b/api/v1/validatorstate_test.go @@ -131,7 +131,10 @@ func TestValidatorStateJSON(t *testing.T) { assert.Equal(t, test.isExited, res.IsExited()) assert.Equal(t, test.hasExited, res.HasExited()) assert.Equal(t, test.hasBalance, res.HasBalance()) - assert.Equal(t, strings.Trim(string(rt), `"`), res.String()) + + valState, err := res.String() + require.NoError(t, err) + assert.Equal(t, strings.Trim(string(rt), `"`), valState) } }) } @@ -261,3 +264,18 @@ func TestValidatorToState(t *testing.T) { }) } } + +func TestString(t *testing.T) { + t.Run("valid state", func(t *testing.T) { + state := api.ValidatorStateActiveOngoing + resp, err := state.String() + require.NoError(t, err) + require.Equal(t, resp, "active_ongoing") + }) + + t.Run("invalid state", func(t *testing.T) { + state := api.ValidatorState(25) + _, err := state.String() + require.Error(t, err) + }) +}