[Internal] Variable storage

sdk/internal/CHANGELOG.md

@@ -1,12 +1,13 @@
 # Release History
 
-## 0.8.3 (2021-11-30)
+## 0.8.3 (2021-12-07)
 
 ### Features Added
 * If `AZURE_RECORD_MODE` is not set, default to `playback`
 * If `PROXY_CERT` is not set, try to find it based on the `GOPATH` environment variable and the path to `eng/common/testproxy/dotnet-devcert.crt`
 * Adds `NewRecordingHTTPClient()` method which returns an `azcore.Transporter` interface that routes requests to the test proxy [#16221](https://github.com/Azure/azure-sdk-for-go/pull/16221).
 * Adds the `SetBodilessMatcher` method [#16256](https://github.com/Azure/azure-sdk-for-go/pull/16256)
+* Added variables storage to the `Stop` function. Pass in a `map[string]interface{}` to the `Stop` method options and the values can be retrieved with the `GetVariables(t *testing.T)` function [#16375](https://github.com/Azure/azure-sdk-for-go/pull/16375).
 
 ### Breaking Changes
 * Renames `ResetSanitizers` to `ResetProxy` [#16256](https://github.com/Azure/azure-sdk-for-go/pull/16256)

sdk/internal/recording/recording.go

@@ -10,6 +10,7 @@ import (
 	"bytes"
 	"crypto/tls"
 	"crypto/x509"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io/ioutil"
@@ -492,6 +493,7 @@ const (
 type recordedTest struct {
 	recordingId string
 	liveOnly    bool
+	variables   map[string]interface{}
 }
 
 var testSuite = map[string]recordedTest{}
@@ -505,6 +507,7 @@ var client = http.Client{
 type RecordingOptions struct {
 	UseHTTPS        bool
 	GroupForReplace string
+	Variables       map[string]interface{}
 }
 
 func defaultOptions() *RecordingOptions {
@@ -548,6 +551,7 @@ func getTestId(pathToRecordings string, t *testing.T) string {
 	return path.Join(pathToRecordings, "recordings", t.Name()+".json")
 }
 
+// Start tells the test proxy to begin accepting requests for a given test
 func Start(t *testing.T, pathToRecordings string, options *RecordingOptions) error {
 	if options == nil {
 		options = defaultOptions()
@@ -580,20 +584,42 @@ func Start(t *testing.T, pathToRecordings string, options *RecordingOptions) err
 	recId := resp.Header.Get(IDHeader)
 	if recId == "" {
 		b, err := ioutil.ReadAll(resp.Body)
+		defer resp.Body.Close()
 		if err != nil {
 			return err
 		}
 		return fmt.Errorf("Recording ID was not returned by the response. Response body: %s", b)
 	}
+
+	// Unmarshal any variables returned by the proxy
+	var m map[string]interface{}
+	body, err := ioutil.ReadAll(resp.Body)
+	defer resp.Body.Close()
+	if err != nil {
+		return err
+	}
+	if len(body) > 0 {
+		err = json.Unmarshal(body, &m)
+		if err != nil {
+			return err
+		}
+	}
+
 	if val, ok := testSuite[t.Name()]; ok {
 		val.recordingId = recId
+		val.variables = m
 		testSuite[t.Name()] = val
 	} else {
-		testSuite[t.Name()] = recordedTest{recordingId: recId, liveOnly: false}
+		testSuite[t.Name()] = recordedTest{
+			recordingId: recId,
+			liveOnly:    false,
+			variables:   m,
+		}
 	}
 	return nil
 }
 
+// Stop tells the test proxy to stop accepting requests for a given test
 func Stop(t *testing.T, options *RecordingOptions) error {
 	if options == nil {
 		options = defaultOptions()
@@ -614,13 +640,32 @@ func Stop(t *testing.T, options *RecordingOptions) error {
 	if err != nil {
 		return err
 	}
+	if len(options.Variables) > 0 {
+		req.Header.Set("Content-Type", "application/json")
+		marshalled, err := json.Marshal(options.Variables)
+		if err != nil {
+			return err
+		}
+		req.Body = ioutil.NopCloser(bytes.NewReader(marshalled))
+		req.ContentLength = int64(len(marshalled))
+	}
+
 	var recTest recordedTest
 	var ok bool
 	if recTest, ok = testSuite[t.Name()]; !ok {
 		return errors.New("Recording ID was never set. Did you call StartRecording?")
 	}
 	req.Header.Set("x-recording-id", recTest.recordingId)
-	_, err = client.Do(req)
+	resp, err := client.Do(req)
+	if resp.StatusCode != 200 {
+		b, err := ioutil.ReadAll(resp.Body)
+		defer resp.Body.Close()
+		if err == nil {
+			return fmt.Errorf("proxy did not stop the recording properly: %s", string(b))
+		}
+		return fmt.Errorf("proxy did not stop the recording properly: %s", err.Error())
+	}
+	_ = resp.Body.Close()
 	return err
 }
 
@@ -687,6 +732,7 @@ func (c RecordingHTTPClient) Do(req *http.Request) (*http.Response, error) {
 	return c.defaultClient.Do(req)
 }
 
+// NewRecordingHTTPClient returns a type that implements `azcore.Transporter`. This will automatically route tests on the `Do` call.
 func NewRecordingHTTPClient(t *testing.T, options *RecordingOptions) (*RecordingHTTPClient, error) {
 	if options == nil {
 		options = &RecordingOptions{UseHTTPS: true}
@@ -721,3 +767,11 @@ func IsLiveOnly(t *testing.T) bool {
 	}
 	return false
 }
+
+// GetVariables returns access to the variables stored by the test proxy for a specific test
+func GetVariables(t *testing.T) map[string]interface{} {
+	if s, ok := testSuite[t.Name()]; ok {
+		return s.variables
+	}
+	return nil
+}

sdk/internal/recording/recording_test.go

@@ -569,22 +569,47 @@ func TestFindProxyCertLocation(t *testing.T) {
 	require.Contains(t, location, filepath.Join("eng", "common", "testproxy", "dotnet-devcert.crt"))
 }
 
-func TestModeNotSet(t *testing.T) {
-	proxyMode, _ := os.LookupEnv("AZURE_RECORD_MODE")
-	defer os.Setenv("AZURE_RECORD_MODE", proxyMode)
+func TestVariables(t *testing.T) {
+	temp := recordMode
+	recordMode = RecordingMode
+	defer func() { recordMode = temp }()
 
-	os.Unsetenv("AZURE_RECORD_MODE")
-	require.Equal(t, PlaybackMode, GetRecordMode())
-}
+	err := Start(t, packagePath, nil)
+	require.NoError(t, err)
+
+	client, err := NewRecordingHTTPClient(t, nil)
+	require.NoError(t, err)
 
-func TestModeNotSetStartStop(t *testing.T) {
-	proxyMode, _ := os.LookupEnv("AZURE_RECORD_MODE")
-	defer os.Setenv("AZURE_RECORD_MODE", proxyMode)
+	req, err := http.NewRequest("POST", "https://azsdkengsys.azurecr.io/acr/v1/some_registry/_tags", nil)
+	require.NoError(t, err)
 
-	os.Unsetenv("AZURE_RECORD_MODE")
-	err := Start(t, packagePath, nil)
+	resp, err := client.Do(req)
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+
+	require.NotNil(t, GetRecordingId(t))
+
+	err = Stop(t, &RecordingOptions{Variables: map[string]interface{}{"key1": "value1", "key2": "1"}})
+	require.NoError(t, err)
+
+	recordMode = PlaybackMode
+	err = Start(t, packagePath, nil)
 	require.NoError(t, err)
-	require.Equal(t, PlaybackMode, GetRecordMode())
+
+	variables := GetVariables(t)
+	require.Equal(t, variables["key1"], "value1")
+	require.Equal(t, variables["key2"], "1")
+
 	err = Stop(t, nil)
 	require.NoError(t, err)
+
+	// Make sure the file is there
+	jsonFile, err := os.Open(fmt.Sprintf("./testdata/recordings/%s.json", t.Name()))
+	require.NoError(t, err)
+	defer func() {
+		err = jsonFile.Close()
+		require.NoError(t, err)
+		err = os.Remove(jsonFile.Name())
+		require.NoError(t, err)
+	}()
 }

sdk/internal/recording/sanitizer.go

@@ -102,6 +102,7 @@ func handleProxyResponse(resp *http.Response, err error) error {
 	}
 
 	body, err := ioutil.ReadAll(resp.Body)
+	defer resp.Body.Close()
 	if err != nil {
 		return err
 	}