From 56672803690350905e8d224477b2c38209e8a57e Mon Sep 17 00:00:00 2001
From: Joshua Spence <josh@spence.com.au>
Date: Fri, 2 Sep 2022 13:03:28 +1000
Subject: [PATCH] Add support for error responses

---
 pkg/mockclient/expectations.go      | 25 +++++++++++++++++++++++++
 pkg/mockclient/expectations_test.go |  9 +++++++++
 2 files changed, 34 insertions(+)

diff --git a/pkg/mockclient/expectations.go b/pkg/mockclient/expectations.go
index 6af13ea..5041138 100644
--- a/pkg/mockclient/expectations.go
+++ b/pkg/mockclient/expectations.go
@@ -2,6 +2,7 @@ package mockclient
 
 import (
 	"fmt"
+	"reflect"
 	"time"
 )
 
@@ -9,6 +10,7 @@ import (
 type Expectation struct {
 	Request  *RequestMatcher `json:"httpRequest"`
 	Response *ActionResponse `json:"httpResponse,omitempty"`
+	Error    *ActionError    `json:"httpError,omitempty"`
 	Times    *Times          `json:"times,omitempty"`
 }
 
@@ -28,6 +30,11 @@ type ActionResponse struct {
 	Delay      *Delay              `json:"delay,omitempty"`
 }
 
+// ActionError defines the failure mode
+type ActionError struct {
+	DropConnection bool `json:"dropConnection,omitempty"`
+}
+
 // ResponseBody defines the request body the MockServer will return when serving a matched response
 type ResponseBody struct {
 	Type   string `json:"type"`
@@ -60,6 +67,7 @@ func CreateExpectation(opts ...ExpectationOption) *Expectation {
 			Path: "/(.*)",
 		},
 		Response: &ActionResponse{},
+		Error:    &ActionError{},
 	}
 
 	// Append all options that are set (discard defaults)
@@ -67,6 +75,14 @@ func CreateExpectation(opts ...ExpectationOption) *Expectation {
 		e = opt(e)
 	}
 
+	if reflect.DeepEqual(e.Response, &ActionResponse{}) {
+		e.Response = nil
+	}
+
+	if reflect.DeepEqual(e.Error, &ActionError{}) {
+		e.Error = nil
+	}
+
 	return e
 }
 
@@ -196,6 +212,15 @@ func ThenResponseDelay(delay time.Duration) ExpectationOption {
 	}
 }
 
+// ThenDropConnection creates an action that prematurely closes the connection
+func ThenDropConnection() ExpectationOption {
+	return func(e *Expectation) *Expectation {
+		r := e.Error
+		r.DropConnection = true
+		return e
+	}
+}
+
 func boolPointer(value bool) *bool {
 	b := value
 	return &b
diff --git a/pkg/mockclient/expectations_test.go b/pkg/mockclient/expectations_test.go
index 04aa606..a6dd2b9 100644
--- a/pkg/mockclient/expectations_test.go
+++ b/pkg/mockclient/expectations_test.go
@@ -131,6 +131,15 @@ func TestExpectations(t *testing.T) {
   				"unlimited" : false
 			}
 		}`},
+		{"Path should be matched and then an error occurs", CreateExpectation(WhenRequestPath("/path"), ThenDropConnection()), `
+    {
+      "httpRequest": {
+        "path": "/path"
+      },
+      "httpError": {
+        "dropConnection": true
+      }
+    }`},
 	}
 
 	for _, tc := range testCases {