feat: allow users to set custom headers for recipes (#2424)

Closes #1971 --------- Co-authored-by: leoporoli <[email protected]>
kurtosis-tech · May 7, 2024 · 07d6ebf · 07d6ebf
1 parent 2439b39
commit 07d6ebf
Show file tree

Hide file tree

Showing 19 changed files with 110 additions and 27 deletions.
diff --git a/core/server/api_container/server/service_network/default_service_network.go b/core/server/api_container/server/service_network/default_service_network.go
@@ -628,7 +628,7 @@ func (network *DefaultServiceNetwork) RunExecs(ctx context.Context, userServiceC
 	return successfulExecs, failedExecs, nil
 }
 
-func (network *DefaultServiceNetwork) HttpRequestService(ctx context.Context, serviceIdentifier string, portId string, method string, contentType string, endpoint string, body string) (*http.Response, error) {
+func (network *DefaultServiceNetwork) HttpRequestService(ctx context.Context, serviceIdentifier string, portId string, method string, contentType string, endpoint string, body string, headers map[string]string) (*http.Response, error) {
 	logrus.Debugf("Making a request '%v' '%v' '%v' '%v' '%v' '%v'", serviceIdentifier, portId, method, contentType, endpoint, body)
 	userService, getServiceErr := network.GetService(ctx, serviceIdentifier)
 	if getServiceErr != nil {
@@ -643,6 +643,9 @@ func (network *DefaultServiceNetwork) HttpRequestService(ctx context.Context, se
 	if err != nil {
 		return nil, stacktrace.NewError("An error occurred building HTTP request on service '%v', URL '%v'", userService, url)
 	}
+	for header, headerValue := range headers {
+		req.Header.Set(header, headerValue)
+	}
 	if contentType != "" {
 		req.Header.Set("Content-Type", contentType)
 	}

diff --git a/core/server/api_container/server/service_network/mock_service_network.go b/core/server/api_container/server/service_network/mock_service_network.go
diff --git a/core/server/api_container/server/service_network/service_network.go b/core/server/api_container/server/service_network/service_network.go
@@ -90,7 +90,7 @@ type ServiceNetwork interface {
 		error,
 	)
 
-	HttpRequestService(ctx context.Context, serviceIdentifier string, portId string, method string, contentType string, endpoint string, body string) (*http.Response, error)
+	HttpRequestService(ctx context.Context, serviceIdentifier string, portId string, method string, contentType string, endpoint string, body string, headers map[string]string) (*http.Response, error)
 
 	GetService(ctx context.Context, serviceIdentifier string) (*service.Service, error)
 

diff --git a/...r/startosis_engine/kurtosis_starlark_framework/test_engine/add_services_framework_test.go b/...r/startosis_engine/kurtosis_starlark_framework/test_engine/add_services_framework_test.go
@@ -127,6 +127,7 @@ func (suite *KurtosisPlanInstructionTestSuite) TestAddServices() {
 		"",
 		testReadyConditionsRecipeEndpoint,
 		"",
+		testEmptyHeaders,
 	).Times(1).Return(&http.Response{
 		Status:     "200 OK",
 		StatusCode: 200,
@@ -155,6 +156,7 @@ func (suite *KurtosisPlanInstructionTestSuite) TestAddServices() {
 		"",
 		testReadyConditions2RecipeEndpoint,
 		"",
+		testEmptyHeaders,
 	).Times(1).Return(&http.Response{
 		Status:     "201 OK",
 		StatusCode: 201,

diff --git a/...gine/kurtosis_starlark_framework/test_engine/get_http_request_recipe_no_extractor_test.go b/...gine/kurtosis_starlark_framework/test_engine/get_http_request_recipe_no_extractor_test.go
@@ -30,6 +30,7 @@ func (suite *KurtosisTypeConstructorTestSuite) TestGetHttpRequestRecipeNoExtract
 		"",
 		"/test",
 		"",
+		testEmptyHeaders,
 	).Times(1).Return(
 		&http.Response{
 			Status:           "200 OK",

diff --git a/.../startosis_engine/kurtosis_starlark_framework/test_engine/get_http_request_recipe_test.go b/.../startosis_engine/kurtosis_starlark_framework/test_engine/get_http_request_recipe_test.go
@@ -34,6 +34,7 @@ func (suite *KurtosisTypeConstructorTestSuite) TestGetHttpRequestRecipe() {
 		"",
 		"/test",
 		"",
+		testEmptyHeaders,
 	).Times(1).Return(
 		&http.Response{
 			Status:           "200 OK",

diff --git a/...s_engine/kurtosis_starlark_framework/test_engine/post_http_request_recipe_minimal_test.go b/...s_engine/kurtosis_starlark_framework/test_engine/post_http_request_recipe_minimal_test.go
@@ -30,6 +30,7 @@ func (suite *KurtosisTypeConstructorTestSuite) TestPostHttpRequestRecipeMinimal(
 		"application/json",
 		"/test",
 		"",
+		testEmptyHeaders,
 	).Times(1).Return(
 		&http.Response{
 			Status:           "200 OK",

diff --git a/...startosis_engine/kurtosis_starlark_framework/test_engine/post_http_request_recipe_test.go b/...startosis_engine/kurtosis_starlark_framework/test_engine/post_http_request_recipe_test.go
@@ -34,6 +34,7 @@ func (suite *KurtosisTypeConstructorTestSuite) TestPostHttpRequestRecipe() {
 		"application/json",
 		"/test",
 		"{}",
+		map[string]string{"key": "value"},
 	).Times(1).Return(
 		&http.Response{
 			Status:           "200 OK",
@@ -63,7 +64,8 @@ func (suite *KurtosisTypeConstructorTestSuite) TestPostHttpRequestRecipe() {
 
 func (t *postHttpRequestRecipeTestCase) GetStarlarkCode() string {
 	extractors := `{"result": ".value"}`
-	return fmt.Sprintf("%s(%s=%q, %s=%q, %s=%q, %s=%q, %s=%s)", recipe.PostHttpRecipeTypeName, recipe.PortIdAttr, testPrivatePortId, recipe.EndpointAttr, "/test", recipe.RequestBodyAttr, "{}", recipe.ContentTypeAttr, "application/json", recipe.ExtractAttr, extractors)
+	headers := `{"key": "value"}`
+	return fmt.Sprintf("%s(%s=%q, %s=%q, %s=%q, %s=%q, %s=%s, %s=%s)", recipe.PostHttpRecipeTypeName, recipe.PortIdAttr, testPrivatePortId, recipe.EndpointAttr, "/test", recipe.RequestBodyAttr, "{}", recipe.ContentTypeAttr, "application/json", recipe.ExtractAttr, extractors, recipe.HeadersAttr, headers)
 }
 
 func (t *postHttpRequestRecipeTestCase) Assert(typeValue builtin_argument.KurtosisValueType) {

diff --git a/...tosis_engine/kurtosis_starlark_framework/test_engine/request_named_args_framework_test.go b/...tosis_engine/kurtosis_starlark_framework/test_engine/request_named_args_framework_test.go
@@ -42,6 +42,7 @@ func (suite *KurtosisPlanInstructionTestSuite) TestRequestWithNamedArgs() {
 		requestContentType,
 		requestEndpoint,
 		requestBody,
+		testEmptyHeaders,
 	).Times(1).Return(
 		&http.Response{
 			Status:           "200 OK",

diff --git a/..._engine/kurtosis_starlark_framework/test_engine/request_positional_args_framework_test.go b/..._engine/kurtosis_starlark_framework/test_engine/request_positional_args_framework_test.go
@@ -30,6 +30,7 @@ func (suite *KurtosisPlanInstructionTestSuite) TestRequestWithPositionalArgs() {
 		requestContentType,
 		requestEndpoint,
 		requestBody,
+		testEmptyHeaders,
 	).Times(1).Return(
 		&http.Response{
 			Status:           "200 OK",

diff --git a/...ainer/server/startosis_engine/kurtosis_starlark_framework/test_engine/static_constants.go b/...ainer/server/startosis_engine/kurtosis_starlark_framework/test_engine/static_constants.go
@@ -115,6 +115,7 @@ var (
 	testReadyConditionsTarget         = "200"
 	testReadyConditionsInterval       = "1s"
 	testReadyConditionsTimeout        = "100ms"
+	testEmptyHeaders                  = map[string]string{}
 
 	testReadyConditions2RecipePortId   = "https"
 	testReadyConditions2RecipeEndpoint = "/user-access"

diff --git a/...tartosis_engine/kurtosis_starlark_framework/test_engine/wait_named_args_framework_test.go b/...tartosis_engine/kurtosis_starlark_framework/test_engine/wait_named_args_framework_test.go
@@ -48,6 +48,7 @@ func (suite *KurtosisPlanInstructionTestSuite) TestWaitWithNamedArgs() {
 		waitRecipeContentType,
 		waitRecipeEndpoint,
 		waitRecipeBody,
+		testEmptyHeaders,
 	).Times(1).Return(
 		&http.Response{
 			Status:           "200 OK",

diff --git a/...sis_engine/kurtosis_starlark_framework/test_engine/wait_positional_args_framework_test.go b/...sis_engine/kurtosis_starlark_framework/test_engine/wait_positional_args_framework_test.go
@@ -32,6 +32,7 @@ func (suite *KurtosisPlanInstructionTestSuite) TestWaitWithPositionalArgs() {
 		waitRecipeContentType,
 		waitRecipeEndpoint,
 		waitRecipeBody,
+		testEmptyHeaders,
 	).Times(1).Return(
 		&http.Response{
 			Status:           "200 OK",

diff --git a/core/server/api_container/server/startosis_engine/recipe/get_http_request_recipe.go b/core/server/api_container/server/startosis_engine/recipe/get_http_request_recipe.go
@@ -52,6 +52,15 @@ func NewGetHttpRequestRecipeType() *kurtosis_type_constructor.KurtosisTypeConstr
 						return interpretationErr
 					},
 				},
+				{
+					Name:              HeadersAttr,
+					IsOptional:        true,
+					ZeroValueProvider: builtin_argument.ZeroValueProvider[*starlark.Dict],
+					Validator: func(value starlark.Value) *startosis_errors.InterpretationError {
+						_, interpretationErr := convertHeadersToMapStringString(true, value)
+						return interpretationErr
+					},
+				},
 			},
 		},
 		Instantiate: instantiateGetHttpRequestRecipe,
@@ -118,6 +127,16 @@ func (recipe *GetHttpRequestRecipe) Execute(
 		return nil, interpretationErr
 	}
 
+	rawHeaders, found, interpretationErr := kurtosis_type_constructor.ExtractAttrValue[*starlark.Dict](
+		recipe.KurtosisValueTypeDefault, HeadersAttr)
+	if interpretationErr != nil {
+		return nil, interpretationErr
+	}
+	headers, interpretationErr := convertHeadersToMapStringString(found, rawHeaders)
+	if interpretationErr != nil {
+		return nil, interpretationErr
+	}
+
 	serviceNameStr := string(serviceName)
 	if serviceNameStr == "" {
 		return nil, stacktrace.NewError("The service name parameter can't be an empty string")
@@ -134,6 +153,7 @@ func (recipe *GetHttpRequestRecipe) Execute(
 		noContentType,
 		endpoint.GoString(),
 		extractors,
+		headers,
 	)
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "An error occurred when running HTTP request recipe '%v'", recipe.String())

diff --git a/core/server/api_container/server/startosis_engine/recipe/http_request_recipe.go b/core/server/api_container/server/startosis_engine/recipe/http_request_recipe.go
@@ -7,6 +7,7 @@ import (
 	"github.com/kurtosis-tech/kurtosis/core/server/api_container/server/service_network"
 	"github.com/kurtosis-tech/kurtosis/core/server/api_container/server/startosis_engine/kurtosis_instruction/shared_helpers/magic_string_helper"
 	"github.com/kurtosis-tech/kurtosis/core/server/api_container/server/startosis_engine/kurtosis_starlark_framework/builtin_argument"
+	"github.com/kurtosis-tech/kurtosis/core/server/api_container/server/startosis_engine/kurtosis_types"
 	"github.com/kurtosis-tech/kurtosis/core/server/api_container/server/startosis_engine/runtime_value_store"
 	"github.com/kurtosis-tech/kurtosis/core/server/api_container/server/startosis_engine/startosis_errors"
 	"github.com/kurtosis-tech/kurtosis/core/server/api_container/server/startosis_engine/startosis_validator"
@@ -27,6 +28,7 @@ const (
 	// Common attributes for both [Get|Post]HttpRequestRecipe
 	PortIdAttr   = "port_id"
 	EndpointAttr = "endpoint"
+	HeadersAttr  = "headers"
 	ExtractAttr  = "extract"
 )
 
@@ -50,6 +52,7 @@ func executeInternal(
 	contentType string,
 	endpoint string,
 	extractors map[string]string,
+	headers map[string]string,
 ) (map[string]starlark.Comparable, error) {
 	var response *http.Response
 	var err error
@@ -63,15 +66,7 @@ func executeInternal(
 		return nil, stacktrace.NewError("The service name parameter can't be an empty string")
 	}
 
-	response, err = serviceNetwork.HttpRequestService(
-		ctx,
-		serviceNameStr,
-		portId,
-		method,
-		contentType,
-		endpoint,
-		recipeBodyWithRuntimeValue,
-	)
+	response, err = serviceNetwork.HttpRequestService(ctx, serviceNameStr, portId, method, contentType, endpoint, recipeBodyWithRuntimeValue, headers)
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "An error occurred when running HTTP request recipe")
 	}
@@ -135,6 +130,21 @@ func createStarlarkReturnValueInternal(resultUuid string, extractors map[string]
 	return dict, nil
 }
 
+func convertHeadersToMapStringString(isSet bool, headersStarlarkValue starlark.Value) (map[string]string, *startosis_errors.InterpretationError) {
+	if !isSet {
+		return map[string]string{}, nil
+	}
+	headersDict, ok := headersStarlarkValue.(*starlark.Dict)
+	if !ok {
+		return nil, startosis_errors.NewInterpretationError("expected '%v' to be a starlark dict but got '%v'", headersStarlarkValue, reflect.TypeOf(headersStarlarkValue))
+	}
+	headers, interpretationErr := kurtosis_types.SafeCastToMapStringString(headersDict, HeadersAttr)
+	if interpretationErr != nil {
+		return nil, interpretationErr
+	}
+	return headers, nil
+}
+
 func convertExtractorsToDict(isAttrSet bool, extractorsValue starlark.Value) (map[string]string, *startosis_errors.InterpretationError) {
 	extractorStringMap := map[string]string{}
 	if !isAttrSet {

diff --git a/core/server/api_container/server/startosis_engine/recipe/post_http_request_recipe.go b/core/server/api_container/server/startosis_engine/recipe/post_http_request_recipe.go
@@ -70,6 +70,15 @@ func NewPostHttpRequestRecipeType() *kurtosis_type_constructor.KurtosisTypeConst
 						return interpretationErr
 					},
 				},
+				{
+					Name:              HeadersAttr,
+					IsOptional:        true,
+					ZeroValueProvider: builtin_argument.ZeroValueProvider[*starlark.Dict],
+					Validator: func(value starlark.Value) *startosis_errors.InterpretationError {
+						_, interpretationErr := convertHeadersToMapStringString(true, value)
+						return interpretationErr
+					},
+				},
 			},
 		},
 		Instantiate: instantiatePostHttpRequestRecipe,
@@ -154,6 +163,16 @@ func (recipe *PostHttpRequestRecipe) Execute(
 		return nil, interpretationErr
 	}
 
+	rawHeaders, found, interpretationErr := kurtosis_type_constructor.ExtractAttrValue[*starlark.Dict](
+		recipe.KurtosisValueTypeDefault, HeadersAttr)
+	if interpretationErr != nil {
+		return nil, interpretationErr
+	}
+	headers, interpretationErr := convertHeadersToMapStringString(found, rawHeaders)
+	if interpretationErr != nil {
+		return nil, interpretationErr
+	}
+
 	serviceNameStr := string(serviceName)
 	if serviceNameStr == "" {
 		return nil, stacktrace.NewError("The service name parameter can't be an empty string")
@@ -170,6 +189,7 @@ func (recipe *PostHttpRequestRecipe) Execute(
 		contentType.GoString(),
 		endpoint.GoString(),
 		extractors,
+		headers,
 	)
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "An error occurred when running HTTP request recipe '%v'", recipe.String())

diff --git a/docs/docs/api-reference/starlark-reference/get-http-request-recipe.md b/docs/docs/api-reference/starlark-reference/get-http-request-recipe.md
@@ -28,6 +28,12 @@ get_request_recipe = GetHttpRequestRecipe(
     extract = {
         "extractfield" : ".name.id",
     },
+
+    # This field allows you to pass custom headers with the GET request
+    # OPTIONAL (Default: {})
+    headers = {
+        "Authorization": "Bearer my.secret.token" 
+    },
 )
 ```
 

diff --git a/docs/docs/api-reference/starlark-reference/post-http-request-recipe.md b/docs/docs/api-reference/starlark-reference/post-http-request-recipe.md
@@ -31,6 +31,14 @@ post_request_recipe = PostHttpRequestRecipe(
     extract = {
         "extractfield" : ".name.id",
     },
+
+    # This field allows you to pass custom headers with the POST request
+    # Any content type passed through the content_type attribtue will hold preference
+    # over content type passed through this
+    # OPTIONAL (Default: {})
+    headers = {
+        "Authorization": "Bearer my.secret.token"
+    },
 )
 ```
 

diff --git a/...estsuite/startosis_request_wait_assert_test/startosis_complex_request_wait_assert_test.go b/...estsuite/startosis_request_wait_assert_test/startosis_complex_request_wait_assert_test.go
@@ -39,16 +39,19 @@ def run(plan):
 		port_id = "http-port",
 		endpoint = "/",
 		content_type="text/plain",
+		headers={"fizz":"buzz"},
 		body=response["extract.exploded-slash"],
 		extract = {
 			"my-body": ".body",
-			"my-content-type": '.headers["content-type"]'
+			"my-content-type": '.headers["content-type"]',
+			"my-headers-fizz": '.headers["fizz"]'
 		}
 	)
 	plan.wait(recipe=post_recipe, field="code", assertion="==", target_value=200, service_name="web-server-complex-request-wait-test")
 	post_response = plan.request(recipe=post_recipe, service_name="web-server-complex-request-wait-test")
 	plan.verify(post_response["code"], "==", 200)
 	plan.verify(post_response["extract.my-content-type"], "==", "text/plain")
+	plan.verify(post_response["extract.my-headers-fizz"], "==", "buzz")
 	plan.verify(post_response["extract.my-body"], "==", "bar")
 	post_recipe_no_body = PostHttpRequestRecipe(
 		port_id = "http-port",
-Original file line number
+Diff line change
@@ Expand Up / @@ -90,7 +90,7 @@ type ServiceNetwork interface { @@
     		error,
     	)
-    	HttpRequestService(ctx context.Context, serviceIdentifier string, portId string, method string, contentType string, endpoint string, body string) (*http.Response, error)
+    	HttpRequestService(ctx context.Context, serviceIdentifier string, portId string, method string, contentType string, endpoint string, body string, headers map[string]string) (*http.Response, error)
     	GetService(ctx context.Context, serviceIdentifier string) (*service.Service, error)
@@ Expand Down @@