header date middleware

yandex · May 18, 2023 · 449b04d · 449b04d
1 parent 19ef8ea
commit 449b04d
Show file tree

Hide file tree

Showing 6 changed files with 333 additions and 27 deletions.
diff --git a/components/phttp/import/import.go b/components/phttp/import/import.go
@@ -11,7 +11,6 @@ import (
 	"github.com/spf13/afero"
 	phttp "github.com/yandex/pandora/components/guns/http"
 	httpProvider "github.com/yandex/pandora/components/providers/http"
-	"github.com/yandex/pandora/components/providers/http/config"
 	"github.com/yandex/pandora/core"
 	"github.com/yandex/pandora/core/register"
 	"github.com/yandex/pandora/lib/answlog"
@@ -20,29 +19,7 @@ import (
 )
 
 func Import(fs afero.Fs) {
-	register.Provider("http", func(conf config.Config) (core.Provider, error) {
-		return httpProvider.NewProvider(fs, conf)
-	})
-
-	register.Provider("http/json", func(conf config.Config) (core.Provider, error) {
-		conf.Decoder = config.DecoderJSONLine
-		return httpProvider.NewProvider(fs, conf)
-	})
-
-	register.Provider("uri", func(conf config.Config) (core.Provider, error) {
-		conf.Decoder = config.DecoderURI
-		return httpProvider.NewProvider(fs, conf)
-	})
-
-	register.Provider("uripost", func(conf config.Config) (core.Provider, error) {
-		conf.Decoder = config.DecoderURIPost
-		return httpProvider.NewProvider(fs, conf)
-	})
-
-	register.Provider("raw", func(conf config.Config) (core.Provider, error) {
-		conf.Decoder = config.DecoderRaw
-		return httpProvider.NewProvider(fs, conf)
-	})
+	httpProvider.Import(fs)
 
 	register.Gun("http", func(conf phttp.HTTPGunConfig) func() core.Gun {
 		targetResolved, _ := PreResolveTargetAddr(&conf.Client, conf.Gun.Target)

diff --git a/components/providers/http/import.go b/components/providers/http/import.go
@@ -0,0 +1,42 @@
+package http
+
+import (
+	"github.com/spf13/afero"
+	"github.com/yandex/pandora/components/providers/http/config"
+	"github.com/yandex/pandora/components/providers/http/middleware"
+	headerdate "github.com/yandex/pandora/components/providers/http/middleware/headerdate"
+	httpRegister "github.com/yandex/pandora/components/providers/http/register"
+	"github.com/yandex/pandora/core"
+	"github.com/yandex/pandora/core/register"
+)
+
+func Import(fs afero.Fs) {
+	register.Provider("http", func(cfg config.Config) (core.Provider, error) {
+		return NewProvider(fs, cfg)
+	})
+
+	register.Provider("http/json", func(cfg config.Config) (core.Provider, error) {
+		cfg.Decoder = config.DecoderJSONLine
+		return NewProvider(fs, cfg)
+	})
+
+	register.Provider("uri", func(cfg config.Config) (core.Provider, error) {
+		cfg.Decoder = config.DecoderURI
+		return NewProvider(fs, cfg)
+	})
+
+	register.Provider("uripost", func(cfg config.Config) (core.Provider, error) {
+		cfg.Decoder = config.DecoderURIPost
+		return NewProvider(fs, cfg)
+	})
+
+	register.Provider("raw", func(cfg config.Config) (core.Provider, error) {
+		cfg.Decoder = config.DecoderRaw
+		return NewProvider(fs, cfg)
+	})
+
+	httpRegister.HTTPMW("header/date", func(cfg headerdate.Config) (middleware.Middleware, error) {
+		return headerdate.NewMiddleware(cfg)
+	})
+
+}
diff --git a/components/providers/http/middleware/headerdate/middleware.go b/components/providers/http/middleware/headerdate/middleware.go
@@ -0,0 +1,47 @@
+package headerdate
+
+import (
+	"context"
+	"net/http"
+	"time"
+
+	"go.uber.org/zap"
+)
+
+const defaultHeaderName = "Date"
+
+type Config struct {
+	Location   string
+	HeaderName string
+}
+
+func NewMiddleware(cfg Config) (*Middleware, error) {
+	m := &Middleware{location: time.UTC, header: defaultHeaderName}
+
+	if cfg.Location != "" {
+		loc, err := time.LoadLocation(cfg.Location)
+		if err != nil {
+			return nil, err
+		}
+		m.location = loc
+	}
+	if cfg.HeaderName != "" {
+		m.header = cfg.HeaderName
+	}
+
+	return m, nil
+}
+
+type Middleware struct {
+	location *time.Location
+	header   string
+}
+
+func (m *Middleware) InitMiddleware(ctx context.Context, log *zap.Logger) error {
+	return nil
+}
+
+func (m *Middleware) UpdateRequest(req *http.Request) error {
+	req.Header.Add(m.header, time.Now().In(m.location).Format(http.TimeFormat))
+	return nil
+}
diff --git a/components/providers/http/middleware/headerdate/middleware_test.go b/components/providers/http/middleware/headerdate/middleware_test.go
@@ -0,0 +1,84 @@
+package headerdate
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/zap"
+)
+
+func TestMiddleware_InitMiddleware(t *testing.T) {
+	cfg := Config{Location: ""}
+	middleware, err := NewMiddleware(cfg)
+	assert.NoError(t, err)
+
+	err = middleware.InitMiddleware(context.Background(), zap.NewNop())
+	assert.NoError(t, err)
+}
+
+func TestMiddleware_UpdateRequest(t *testing.T) {
+	t.Run("empty config", func(t *testing.T) {
+		cfg := Config{}
+		middleware, err := NewMiddleware(cfg)
+		assert.NoError(t, err)
+
+		req := httptest.NewRequest(http.MethodGet, "/", nil)
+
+		err = middleware.UpdateRequest(req)
+		assert.NoError(t, err)
+
+		dateHeader := req.Header.Get("Date")
+		assert.NotEmpty(t, dateHeader)
+
+		expectedDate := time.Now().In(time.UTC).Format(http.TimeFormat)
+		assert.Equal(t, expectedDate, dateHeader)
+	})
+	t.Run("America/New_York", func(t *testing.T) {
+		loc, err := time.LoadLocation("America/New_York")
+		assert.NoError(t, err)
+
+		cfg := Config{Location: "America/New_York"}
+		middleware, err := NewMiddleware(cfg)
+		assert.NoError(t, err)
+
+		req := httptest.NewRequest(http.MethodGet, "/", nil)
+
+		err = middleware.UpdateRequest(req)
+		assert.NoError(t, err)
+
+		dateHeader := req.Header.Get("Date")
+		assert.NotEmpty(t, dateHeader)
+
+		expectedDate := time.Now().In(loc).Format(http.TimeFormat)
+		assert.Equal(t, expectedDate, dateHeader)
+	})
+	t.Run("custom header name", func(t *testing.T) {
+		loc, err := time.LoadLocation("America/New_York")
+		assert.NoError(t, err)
+
+		cfg := Config{Location: "America/New_York", HeaderName: "CreatedDate"}
+		middleware, err := NewMiddleware(cfg)
+		assert.NoError(t, err)
+
+		req := httptest.NewRequest(http.MethodGet, "/", nil)
+
+		err = middleware.UpdateRequest(req)
+		assert.NoError(t, err)
+
+		dateHeader := req.Header.Get("CreatedDate")
+		assert.NotEmpty(t, dateHeader)
+
+		expectedDate := time.Now().In(loc).Format(http.TimeFormat)
+		assert.Equal(t, expectedDate, dateHeader)
+	})
+}
+
+func TestMiddleware_UpdateRequest_InvalidLocation(t *testing.T) {
+	cfg := Config{Location: "Invalid/Location"}
+	_, err := NewMiddleware(cfg)
+	assert.Error(t, err)
+}
diff --git a/components/providers/http/provider_test.go b/components/providers/http/provider_test.go
@@ -0,0 +1,118 @@
+package http
+
+import (
+	"os"
+	"testing"
+
+	"github.com/spf13/afero"
+	"github.com/yandex/pandora/components/providers/http/config"
+	"github.com/yandex/pandora/components/providers/http/provider"
+)
+
+func TestNewProvider_invalidDecoder(t *testing.T) {
+	fs := afero.NewMemMapFs()
+	conf := config.Config{
+		Decoder: "invalid",
+	}
+
+	p, err := NewProvider(fs, conf)
+	if p != nil || err == nil {
+		t.Error("expected error when creating provider with invalid decoder type")
+	}
+}
+
+func TestNewProvider(t *testing.T) {
+	fs := afero.NewMemMapFs()
+
+	t.Run("InvalidDecoder", func(t *testing.T) {
+	})
+
+	tmpFile, err := fs.Create("ammo")
+	if err != nil {
+		t.Fatalf("failed to create temp file: %s", err)
+	}
+	defer os.Remove(tmpFile.Name())
+
+	if _, err := tmpFile.Write([]byte("GET / HTTP/1.1\nHost: example.com\n\n")); err != nil {
+		t.Fatalf("failed to write data to temp file: %s", err)
+	}
+
+	cases := []struct {
+		name     string
+		conf     config.Config
+		expected config.DecoderType
+		filePath string
+	}{
+		{
+			name:     "URIDecoder",
+			conf:     config.Config{Decoder: config.DecoderURI, Uris: []string{"http://example.com"}},
+			expected: config.DecoderURI,
+			filePath: "",
+		},
+		{
+			name:     "FileDecoder",
+			conf:     config.Config{Decoder: config.DecoderURI, File: tmpFile.Name()},
+			expected: config.DecoderURI,
+			filePath: tmpFile.Name(),
+		},
+		{
+			name:     "DecoderURIPost",
+			conf:     config.Config{Decoder: config.DecoderURIPost, File: tmpFile.Name()},
+			expected: config.DecoderURIPost,
+			filePath: tmpFile.Name(),
+		},
+		{
+			name:     "DecoderRaw",
+			conf:     config.Config{Decoder: config.DecoderRaw, File: tmpFile.Name()},
+			expected: config.DecoderRaw,
+			filePath: tmpFile.Name(),
+		},
+		{
+			name:     "DecoderJSONLine",
+			conf:     config.Config{Decoder: config.DecoderJSONLine, File: tmpFile.Name()},
+			expected: config.DecoderJSONLine,
+			filePath: tmpFile.Name(),
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			providr, err := NewProvider(fs, tc.conf)
+			if err != nil {
+				t.Fatalf("failed to create provider: %s", err)
+			}
+
+			provider, _ := providr.(*provider.Provider)
+
+			defer func() {
+				if err := provider.Close(); err != nil {
+					t.Fatalf("failed to close provider: %s", err)
+				}
+			}()
+
+			if provider == nil {
+				t.Fatal("provider is nil")
+			}
+
+			if provider.Config.Decoder != tc.expected {
+				t.Errorf("unexpected decoder type: got %s, want %s", provider.Config.Decoder, tc.expected)
+			}
+			if provider.Config.File != tc.filePath {
+				t.Errorf("unexpected file path: got %s, want %s", provider.Config.File, tc.filePath)
+			}
+
+			if provider.Decoder == nil && tc.expected != "" {
+				t.Error("decoder is nil")
+			}
+
+			if provider.FS != fs {
+				t.Errorf("unexpected FS: got %v, want %v", provider.FS, fs)
+			}
+
+			if provider.Sink == nil {
+				t.Error("sink is nil")
+			}
+		})
+	}
+
+}
diff --git a/docs/advanced.rst b/docs/advanced.rst
@@ -1,7 +1,7 @@
-Ammo providers
+HTTP Ammo providers
 ==============
 
-Ammo provider is a source of test data: it makes ammo object.
+HTTP Ammo provider is a source of test data: it makes ammo object.
 
 There is a common rule for any (built-in) provider: data supplied by ammo provider are records that will be pushed via established connection to external host (defined in pandora config via `pool.gun.target` option). Thus, you cannot define in the ammofile to which `physical` host your ammo will be sent.
 
@@ -129,6 +129,7 @@ Ammo filters
 ------------
 
 Each http ammo provider lets you choose specific ammo for your test from ammo file with `chosencases` setting:
+
 .. code-block:: yaml
 
   pools:
@@ -152,4 +153,41 @@ uri-style:
 ::
   /?drg tag1
   /
-  /buy tag2
+  /buy tag2
+
+HTTP Ammo middlewares
+---------------------
+
+HTTP Ammo providers have the ability to modify HTTP request just before execution.
+Middlewares are used for this purpose. An example of Middleware that sets the Date header in a request.
+
+
+.. code-block:: yaml
+
+  pools:
+    - ammo:
+        type: uri
+        ...
+        middlewares:
+          - type: header/date
+            location: EST
+            headerName: Date
+
+List of built-in HTTP Ammo middleware:
+- header/date
+
+You can create your own middleware. But in order to do that you need to register them in custom pandora
+
+.. code-block:: go
+
+  import (
+    "github.com/yandex/pandora/components/providers/http/middleware"
+    "github.com/yandex/pandora/components/providers/http/middleware/headerdate"
+    httpRegister "github.com/yandex/pandora/components/providers/http/register"
+  )
+
+  httpRegister.HTTPMW("header/date", func(cfg headerdate.Config) (middleware.Middleware, error) {
+      return headerdate.NewMiddleware(cfg)
+  })
+
+For more on how to write custom pandora, see :ref:`custom`