add fix and tests

attestantio · Oct 19, 2023 · 2f65e8b · 2f65e8b
1 parent bbf582d
commit 2f65e8b
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 3 deletions.
diff --git 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
@@ -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)
+	})
+}