Allow accepting additional StatusCodes with custom fallback response

docs/data-sources/http.md

@@ -54,6 +54,8 @@ data "http" "example" {
 
 ### Optional
 
+- `fallback_on_status` (Set of Number) A set of integers representing HTTP status codes that are acceptable responses, and may substitute the `fallback_response` if received.
+- `fallback_response` (String) A fallback response value to use as `body` if the response status is a code specified in `fallback_on_status`. If not defined, the actual response body is still returned.
 - `request_headers` (Map of String) A map of request header field names and values.
 
 ### Read-Only

internal/provider/data_source.go

@@ -49,6 +49,24 @@ your control should be treated as untrustworthy.`,
 				},
 			},
 
+			"fallback_on_status": {
+				Description: "A set of integers representing HTTP status codes that are acceptable responses, and may substitute the `fallback_response` if received.",
+				Type:        schema.TypeSet,
+				Optional:    true,
+				Elem: &schema.Schema{
+					Type: schema.TypeInt,
+				},
+			},
+
+			"fallback_response": {
+				Description: "A fallback response value to use as `body` if the response status is a code specified in `fallback_on_status`. If not defined, the actual response body is still returned.",
+				Type:        schema.TypeString,
+				Optional:    true,
+				Elem: &schema.Schema{
+					Type: schema.TypeString,
+				},
+			},
+
 			"body": {
 				Description: "The response body returned as a string.",
 				Type:        schema.TypeString,
@@ -74,6 +92,8 @@ your control should be treated as untrustworthy.`,
 func dataSourceRead(ctx context.Context, d *schema.ResourceData, meta interface{}) (diags diag.Diagnostics) {
 	url := d.Get("url").(string)
 	headers := d.Get("request_headers").(map[string]interface{})
+	fallbackStatusCodes := d.Get("fallback_on_status").(*schema.Set)
+	fallbackResponse, useFallback := d.GetOk("fallback_response")
 
 	client := &http.Client{}
 
@@ -93,7 +113,8 @@ func dataSourceRead(ctx context.Context, d *schema.ResourceData, meta interface{
 
 	defer resp.Body.Close()
 
-	if resp.StatusCode != 200 {
+	statusAcceptable := fallbackStatusCodes.Contains(resp.StatusCode)
+	if resp.StatusCode != 200 && !statusAcceptable {
 		return append(diags, diag.Errorf("HTTP request error. Response code: %d", resp.StatusCode)...)
 	}
 
@@ -106,9 +127,13 @@ func dataSourceRead(ctx context.Context, d *schema.ResourceData, meta interface{
 		})
 	}
 
-	bytes, err := ioutil.ReadAll(resp.Body)
-	if err != nil {
-		return append(diags, diag.FromErr(err)...)
+	if statusAcceptable && useFallback {
+		d.Set("body", fallbackResponse)
+	} else {
+		bytes, err := ioutil.ReadAll(resp.Body)
+		if err = d.Set("body", string(bytes)); err != nil {
+			return append(diags, diag.Errorf("Error setting HTTP response body: %s", err)...)
+		}
 	}
 
 	responseHeaders := make(map[string]string)
@@ -118,10 +143,6 @@ func dataSourceRead(ctx context.Context, d *schema.ResourceData, meta interface{
 		responseHeaders[k] = strings.Join(v, ", ")
 	}
 
-	if err = d.Set("body", string(bytes)); err != nil {
-		return append(diags, diag.Errorf("Error setting HTTP response body: %s", err)...)
-	}
-
 	if err = d.Set("response_headers", responseHeaders); err != nil {
 		return append(diags, diag.Errorf("Error setting HTTP response headers: %s", err)...)
 	}

internal/provider/data_source_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
 )
 
 func TestDataSource_http200(t *testing.T) {
@@ -46,6 +47,57 @@ func TestDataSource_http404(t *testing.T) {
 	})
 }
 
+const testDataSourceConfigWithFallback = `
+data "http" "http_test" {
+  url = "%s/meta_%d.txt"
+  fallback_on_status = toset([%[2]d])
+  fallback_response = "fallback"
+}
+`
+
+func TestDataSource_fallback404(t *testing.T) {
+	testHttpMock := setUpMockHttpServer()
+
+	defer testHttpMock.server.Close()
+
+	resource.UnitTest(t, resource.TestCase{
+		ProviderFactories: testProviders(),
+		Steps: []resource.TestStep{
+			{
+				Config: fmt.Sprintf(testDataSourceConfigWithFallback, testHttpMock.server.URL, 404),
+				Check: resource.ComposeTestCheckFunc(
+					resource.TestCheckResourceAttr("data.http.http_test", "body", "fallback"),
+				),
+			},
+		},
+	})
+}
+
+const testDataSourceConfigWithFallbackStatus = `
+data "http" "http_test" {
+  url = "%s/meta_%d.txt"
+  fallback_on_status = toset([%[2]d])
+}
+`
+
+func TestDataSource_allow404(t *testing.T) {
+	testHttpMock := setUpMockHttpServer()
+
+	defer testHttpMock.server.Close()
+
+	resource.UnitTest(t, resource.TestCase{
+		ProviderFactories: testProviders(),
+		Steps: []resource.TestStep{
+			{
+				Config: fmt.Sprintf(testDataSourceConfigWithFallbackStatus, testHttpMock.server.URL, 404),
+				Check: resource.ComposeTestCheckFunc(
+					resource.TestCheckResourceAttr("data.http.http_test", "body", ""),
+				),
+			},
+		},
+	})
+}
+
 func TestDataSource_withHeaders200(t *testing.T) {
 	testHttpMock := setUpMockHttpServer()