GH-85: custom JSON provider

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 {

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)
 }

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()