diff --git a/pkg/smartcontract/callflag/call_flags.go b/pkg/smartcontract/callflag/call_flags.go index c31b819b4..cad3994d8 100644 --- a/pkg/smartcontract/callflag/call_flags.go +++ b/pkg/smartcontract/callflag/call_flags.go @@ -117,3 +117,21 @@ func (f *CallFlag) UnmarshalJSON(data []byte) error { *f = flag return nil } + +// MarshalYAML implements the YAML marshaler interface. +func (f CallFlag) MarshalYAML() (interface{}, error) { + return f.String(), nil +} + +// UnmarshalYAML implements the YAML unmarshaler interface. +func (f *CallFlag) UnmarshalYAML(unmarshal func(interface{}) error) error { + var s string + + err := unmarshal(&s) + if err != nil { + return err + } + + *f, err = FromString(s) + return err +} diff --git a/pkg/smartcontract/callflag/call_flags_test.go b/pkg/smartcontract/callflag/call_flags_test.go index b75145150..b1c057e4e 100644 --- a/pkg/smartcontract/callflag/call_flags_test.go +++ b/pkg/smartcontract/callflag/call_flags_test.go @@ -5,6 +5,7 @@ import ( "github.com/nspcc-dev/neo-go/internal/testserdes" "github.com/stretchr/testify/require" + "gopkg.in/yaml.v2" ) func TestCallFlag_Has(t *testing.T) { @@ -75,3 +76,17 @@ func TestMarshalUnmarshalJSON(t *testing.T) { require.Error(t, err) require.Equal(t, forig, f) } + +func TestMarshalUnmarshalYAML(t *testing.T) { + for _, expected := range []CallFlag{States, States | AllowNotify} { + data, err := yaml.Marshal(expected) + require.NoError(t, err) + + var f CallFlag + require.NoError(t, yaml.Unmarshal(data, &f)) + require.Equal(t, expected, f) + } + + var f CallFlag + require.Error(t, yaml.Unmarshal([]byte(`[]`), &f)) +}