Add chosencases setting (similar to phantom)

add chosencases config option
yandex · Sep 8, 2022 · deb7d99 · deb7d99
1 parent 89d8339
commit deb7d99
Show file tree

Hide file tree

Showing 11 changed files with 114 additions and 13 deletions.
diff --git a/components/grpc/ammo/grpcjson/provider.go b/components/grpc/ammo/grpcjson/provider.go
@@ -13,6 +13,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/spf13/afero"
 	"github.com/yandex/pandora/components/grpc/ammo"
+	"github.com/yandex/pandora/lib/confutil"
 	"go.uber.org/zap"
 )
 
@@ -49,6 +50,7 @@ type Config struct {
 	//Maximum number of byte in an ammo. Default is bufio.MaxScanTokenSize
 	MaxAmmoSize int
 	Source      Source `config:"source"`
+	ChosenCases []string
 }
 
 func (p *Provider) start(ctx context.Context, ammoFile afero.File) error {
@@ -70,6 +72,9 @@ func (p *Provider) start(ctx context.Context, ammoFile afero.File) error {
 					return errors.Wrapf(err, "failed to decode ammo at line: %v; data: %q", line, data)
 				}
 			}
+			if !confutil.IsChosenCase(a.Tag, p.Config.ChosenCases) {
+				continue
+			}
 			ammoNum++
 			select {
 			case p.Sink <- a:

diff --git a/components/phttp/ammo/simple/ammo.go b/components/phttp/ammo/simple/ammo.go
@@ -51,4 +51,8 @@ func (a *Ammo) IsValid() bool {
 	return !a.isInvalid
 }
 
+func (a *Ammo) Tag() string {
+	return a.tag
+}
+
 var _ phttp.Ammo = (*Ammo)(nil)
diff --git a/components/phttp/ammo/simple/jsonline/provider.go b/components/phttp/ammo/simple/jsonline/provider.go
@@ -14,6 +14,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/spf13/afero"
 	"github.com/yandex/pandora/components/phttp/ammo/simple"
+	"github.com/yandex/pandora/lib/confutil"
 	"go.uber.org/zap"
 )
 
@@ -41,10 +42,12 @@ type Config struct {
 	ContinueOnError bool
 	//Maximum number of byte in an ammo. Default is bufio.MaxScanTokenSize
 	MaxAmmoSize int
+	ChosenCases []string
 }
 
 func (p *Provider) start(ctx context.Context, ammoFile afero.File) error {
 	var ammoNum, passNum int
+
 	for {
 		passNum++
 		scanner := bufio.NewScanner(ammoFile)
@@ -62,6 +65,9 @@ func (p *Provider) start(ctx context.Context, ammoFile afero.File) error {
 					return errors.Wrapf(err, "failed to decode ammo at line: %v; data: %q", line, data)
 				}
 			}
+			if !confutil.IsChosenCase(a.Tag(), p.Config.ChosenCases) {
+				continue
+			}
 			ammoNum++
 			select {
 			case p.Sink <- a:

diff --git a/components/phttp/ammo/simple/raw/provider.go b/components/phttp/ammo/simple/raw/provider.go
@@ -8,6 +8,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/spf13/afero"
 	"github.com/yandex/pandora/components/phttp/ammo/simple"
+	"github.com/yandex/pandora/lib/confutil"
 	"go.uber.org/zap"
 )
 
@@ -53,7 +54,8 @@ type Config struct {
 	// Redefine HTTP headers
 	Headers []string
 	// Passes limits ammo file passes. Unlimited if zero.
-	Passes int `validate:"min=0"`
+	Passes      int `validate:"min=0"`
+	ChosenCases []string
 }
 
 func NewProvider(fs afero.Fs, conf Config) *Provider {
@@ -100,6 +102,9 @@ func (p *Provider) start(ctx context.Context, ammoFile afero.File) error {
 			if reqSize == 0 {
 				break // start over from the beginning of file if ammo size is 0
 			}
+			if !confutil.IsChosenCase(tag, p.Config.ChosenCases) {
+				continue
+			}
 			buff := make([]byte, reqSize)
 			if n, err := io.ReadFull(reader, buff); err != nil {
 				return errors.Wrapf(err, "failed to read ammo at position: %v; tried to read: %v; have read: %v", filePosition(ammoFile), reqSize, n)

diff --git a/components/phttp/ammo/simple/uri/decoder.go b/components/phttp/ammo/simple/uri/decoder.go
@@ -14,6 +14,7 @@ import (
 
 	"github.com/pkg/errors"
 	"github.com/yandex/pandora/components/phttp/ammo/simple"
+	"github.com/yandex/pandora/lib/confutil"
 )
 
 type decoder struct {
@@ -24,14 +25,16 @@ type decoder struct {
 	ammoNum       int
 	header        http.Header
 	configHeaders []simple.Header
+	chosenCases   []string
 }
 
-func newDecoder(ctx context.Context, sink chan<- *simple.Ammo, pool *sync.Pool) *decoder {
+func newDecoder(ctx context.Context, sink chan<- *simple.Ammo, pool *sync.Pool, chosenCases []string) *decoder {
 	return &decoder{
-		sink:   sink,
-		header: http.Header{},
-		pool:   pool,
-		ctx:    ctx,
+		sink:        sink,
+		header:      http.Header{},
+		pool:        pool,
+		ctx:         ctx,
+		chosenCases: chosenCases,
 	}
 }
 
@@ -56,6 +59,9 @@ func (d *decoder) decodeURI(line []byte) error {
 	if len(parts) > 1 {
 		tag = parts[1]
 	}
+	if !confutil.IsChosenCase(tag, d.chosenCases) {
+		return nil
+	}
 	req, err := http.NewRequest("GET", string(url), nil)
 	if err != nil {
 		return errors.Wrap(err, "uri decode")

diff --git a/components/phttp/ammo/simple/uri/decoder_test.go b/components/phttp/ammo/simple/uri/decoder_test.go
@@ -21,7 +21,7 @@ func newAmmoPool() *sync.Pool {
 var _ = Describe("Decoder", func() {
 	It("uri decode ctx cancel", func() {
 		ctx, cancel := context.WithCancel(context.Background())
-		decoder := newDecoder(ctx, make(chan *simple.Ammo), newAmmoPool())
+		decoder := newDecoder(ctx, make(chan *simple.Ammo), newAmmoPool(), nil)
 		cancel()
 		err := decoder.Decode([]byte("/some/path"))
 		Expect(err).To(Equal(context.Canceled))
@@ -32,7 +32,7 @@ var _ = Describe("Decoder", func() {
 	)
 	BeforeEach(func() {
 		ammoCh = make(chan *simple.Ammo, 10)
-		decoder = newDecoder(context.Background(), ammoCh, newAmmoPool())
+		decoder = newDecoder(context.Background(), ammoCh, newAmmoPool(), nil)
 	})
 	DescribeTable("invalid input",
 		func(line string) {

diff --git a/components/phttp/ammo/simple/uri/provider.go b/components/phttp/ammo/simple/uri/provider.go
@@ -23,8 +23,9 @@ type Config struct {
 	// Additional HTTP headers
 	Headers []string
 	// Passes limits ammo file passes. Unlimited if zero.
-	Passes int `validate:"min=0"`
-	Uris   []string
+	Passes      int `validate:"min=0"`
+	Uris        []string
+	ChosenCases []string
 }
 
 // TODO: pass logger and metricsRegistry
@@ -73,7 +74,7 @@ type Provider struct {
 }
 
 func (p *Provider) start(ctx context.Context, ammoFile afero.File) error {
-	p.decoder = newDecoder(ctx, p.Sink, &p.Pool)
+	p.decoder = newDecoder(ctx, p.Sink, &p.Pool, p.Config.ChosenCases)
 	// parse and prepare Headers from config
 	decodedConfigHeaders, err := simple.DecodeHTTPConfigHeaders(p.Config.Headers)
 	if err != nil {

diff --git a/components/phttp/ammo/simple/uripost/provider.go b/components/phttp/ammo/simple/uripost/provider.go
@@ -11,6 +11,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/spf13/afero"
 	"github.com/yandex/pandora/components/phttp/ammo/simple"
+	"github.com/yandex/pandora/lib/confutil"
 	"go.uber.org/zap"
 )
 
@@ -26,7 +27,8 @@ type Config struct {
 	// Additional HTTP headers
 	Headers []string
 	// Passes limits ammo file passes. Unlimited if zero.
-	Passes int `validate:"min=0"`
+	Passes      int `validate:"min=0"`
+	ChosenCases []string
 }
 
 func NewProvider(fs afero.Fs, conf Config) *Provider {
@@ -90,6 +92,9 @@ func (p *Provider) start(ctx context.Context, ammoFile afero.File) error {
 			if bodySize == 0 {
 				break // start over from the beginning of file if ammo size is 0
 			}
+			if !confutil.IsChosenCase(tag, p.Config.ChosenCases) {
+				continue
+			}
 			buff := make([]byte, bodySize)
 			if n, err := io.ReadFull(reader, buff); err != nil {
 				return errors.Wrapf(err, "failed to read ammo at position: %v; tried to read: %v; have read: %v", filePosition(ammoFile), bodySize, n)

diff --git a/docs/advanced.rst b/docs/advanced.rst
@@ -18,7 +18,7 @@ Pay attention to special header `Host` defined ``outside`` of Headers dictionary
 Ammofile sample:
 ::
 
-  {"uri": "/", "method": "GET", "headers": {"Accept": "*/*", "Accept-Encoding": "gzip, deflate", "User-Agent": "Pandora"}, "host": "example.com"}
+  {"tag": "tag1", "uri": "/", "method": "GET", "headers": {"Accept": "*/*", "Accept-Encoding": "gzip, deflate", "User-Agent": "Pandora"}, "host": "example.com"}
 
 Config sample:
 
@@ -124,3 +124,32 @@ Example:
         headers:
           - "[Host: yourhost.tld]"
           - "[User-Agent: some user agent]"
+
+Ammo filters
+------------
+
+Each http ammo provider lets you choose specific ammo for your test from ammo file with `chosencases` setting:
+.. code-block:: yaml
+
+  pools:
+    - ammo:
+        type: uri                      # ammo format
+        chosencases: ["tag1", "tag2"]  # use only "tag1" and "tag2" ammo for this test
+        file: ./ammofile               # ammo file path
+
+Tags are defined in ammo files as shown below:
+
+http/json:
+::
+  {"tag": "tag1", "uri": "/",
+
+raw (request-style):
+::
+  73 tag1
+  GET / HTTP/1.0
+
+uri-style:
+::
+  /?drg tag1
+  /
+  /buy tag2
diff --git a/lib/confutil/chosen_cases_filter.go b/lib/confutil/chosen_cases_filter.go
@@ -0,0 +1,15 @@
+package confutil
+
+// Creates filter that returns true if ammo tag is in chosenCases. If no chosenCases provided - returns true
+func IsChosenCase(checkCase string, chosenCases []string) bool {
+	if len(chosenCases) == 0 {
+		return true
+	}
+
+	for _, c := range chosenCases {
+		if c == checkCase {
+			return true
+		}
+	}
+	return false
+}
diff --git a/lib/confutil/chosen_cases_filter_test.go b/lib/confutil/chosen_cases_filter_test.go
@@ -0,0 +1,25 @@
+package confutil
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type testCase struct {
+	ammoTag  string
+	expected bool
+}
+
+func TestChosenCases(t *testing.T) {
+	cases := []string{"tag1", "tag3"}
+
+	tests := []testCase{
+		{"tag1", true},
+		{"tag2", false},
+	}
+
+	for _, tc := range tests {
+		assert.Equal(t, tc.expected, IsChosenCase(tc.ammoTag, cases))
+	}
+}