diff --git a/core/provider/decoder.go b/core/provider/decoder.go index 8edb5e3eb..f6c8735c4 100644 --- a/core/provider/decoder.go +++ b/core/provider/decoder.go @@ -31,6 +31,10 @@ type AmmoDecoder interface { Decode(ammo core.Ammo) error } +type AmmoDecoderFunc func(ammo core.Ammo) error + +func (f AmmoDecoderFunc) Decode(ammo core.Ammo) error { return f(ammo) } + // TODO(skipor): test func NewDecodeProvider(newAmmo func() core.Ammo, newDecoder NewAmmoDecoder, conf DecodeProviderConfig) *DecodeProvider { diff --git a/core/provider/json.go b/core/provider/json.go index 1472750a6..d1956580b 100644 --- a/core/provider/json.go +++ b/core/provider/json.go @@ -15,9 +15,23 @@ import ( "github.com/yandex/pandora/lib/ioutil2" ) +// NewJSONProvider returns generic core.Provider that reads JSON data from source and decodes it +// into ammo returned from newAmmo. func NewJSONProvider(newAmmo func() core.Ammo, conf JSONProviderConfig) core.Provider { + return NewCustomJSONProvider(nil, newAmmo, conf) +} + +// NewCustomJSONProvider is like NewJSONProvider, but also allows to wrap JSON decoder, to +// decode data into intermediate struct, but then transform in into desired ammo. +// For example, decode {"body":"some data"} into struct { Data string }, and transform it to +// http.Request. +func NewCustomJSONProvider(wrapDecoder func(decoder AmmoDecoder) AmmoDecoder, newAmmo func() core.Ammo, conf JSONProviderConfig) core.Provider { var newDecoder NewAmmoDecoder = func(deps core.ProviderDeps, source io.Reader) (AmmoDecoder, error) { - return NewJSONAmmoDecoder(source, conf.Buffer.BufferSizeOrDefault()), nil + decoder := NewJSONAmmoDecoder(source, conf.Buffer.BufferSizeOrDefault()) + if wrapDecoder != nil { + decoder = wrapDecoder(decoder) + } + return decoder, nil } return NewDecodeProvider(newAmmo, newDecoder, conf.Decode) } diff --git a/core/provider/json_test.go b/core/provider/json_test.go index 9a9ec5a1e..42f5c1a24 100644 --- a/core/provider/json_test.go +++ b/core/provider/json_test.go @@ -55,7 +55,39 @@ func TestDecodeProviderPasses(t *testing.T) { assert.False(t, ok) } -func TestDecodeProvider(t *testing.T) { +func TestCustomJSONProvider(t *testing.T) { + input := strings.NewReader(` {"data":"first"}`) + conf := DefaultJSONProviderConfig() + conf.Decode.Source = datasource.NewReader(input) + conf.Decode.Limit = 1 + newAmmo := func() core.Ammo { + return &testJSONAmmo{} + } + wrapDecoder := func(decoder AmmoDecoder) AmmoDecoder { + return AmmoDecoderFunc(func(ammo core.Ammo) error { + err := decoder.Decode(ammo) + if err != nil { + return err + } + ammo.(*testJSONAmmo).Data += " transformed" + return nil + }) + } + provider := NewCustomJSONProvider(wrapDecoder, newAmmo, conf) + err := provider.Run(context.Background(), testDeps()) + require.NoError(t, err) + expected := func(data string) *testJSONAmmo { + return &testJSONAmmo{Data: data} + } + ammo, ok := provider.Acquire() + require.True(t, ok) + assert.Equal(t, expected("first transformed"), ammo) + + _, ok = provider.Acquire() + assert.False(t, ok) +} + +func TestJSONProvider(t *testing.T) { input := strings.NewReader(` {"data":"first"} {"data":"second"} `) conf := DefaultJSONProviderConfig()