From 8171f767357b14030ed7670b8940d630893071b8 Mon Sep 17 00:00:00 2001
From: Timo van Veenendaal <timov@microsoft.com>
Date: Mon, 4 Jul 2022 15:16:25 -0700
Subject: [PATCH] Add support for skip mode

---
 sdk/test-utils/recorder/src/recorder.ts | 44 +++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/sdk/test-utils/recorder/src/recorder.ts b/sdk/test-utils/recorder/src/recorder.ts
index 94294c96fe07..a6eb31be39e8 100644
--- a/sdk/test-utils/recorder/src/recorder.ts
+++ b/sdk/test-utils/recorder/src/recorder.ts
@@ -43,6 +43,21 @@ import { isNode } from "@azure/core-util";
 import { env } from "./utils/env";
 import { decodeBase64 } from "./utils/encoding";
 
+/**
+ * Skip mode to be passed to Recorder#setSkipMode.
+ *
+ * There are 3 possible modes:
+ * - "skip-recording": Skips recording all requests and responses until the skip mode is changed.
+ * - "skip-request-body": Skips recording the request body of requests until the skip mode is changed.
+ * - "default": Record all requests and responses as normal.
+ */
+export type SkipMode = "skip-recording" | "skip-request-body" | "default";
+
+const skipModeToHeaderValue: { [k in SkipMode]?: string } = {
+  "skip-recording": "request-response",
+  "skip-request-body": "request-body",
+};
+
 /**
  * This client manages the recorder life cycle and interacts with the proxy-tool to do the recording,
  * eventually save them in record mode and playing them back in playback mode.
@@ -61,6 +76,7 @@ export class Recorder {
   private httpClient?: HttpClient;
   private sessionFile?: string;
   private variables: Record<string, string>;
+  private skipMode: SkipMode = "default";
 
   constructor(private testContext?: Test | undefined) {
     logger.info(`[Recorder#constructor] Creating a recorder instance in ${getTestMode()} mode`);
@@ -122,6 +138,22 @@ export class Recorder {
       request.headers.set("x-recording-id", this.recordingId);
       request.headers.set("x-recording-mode", getTestMode());
 
+      const skipHeaderValue = skipModeToHeaderValue[this.skipMode];
+      if (skipHeaderValue) {
+        logger.verbose(
+          `[Recorder#redirectRequest] Setting x-recording-skip header to ${skipHeaderValue}`
+        );
+        request.headers.set("x-recording-skip", skipHeaderValue);
+      }
+
+      // Drop request body in playback mode if we skipped recording it
+      if (isPlaybackMode() && this.skipMode === "skip-request-body") {
+        logger.verbose(
+          "[Recorder#redirectRequest] Setting request body to null for playback of recording with skipped body"
+        );
+        request.body = null;
+      }
+
       redirectedUrl.host = testProxyUrl.host;
       redirectedUrl.port = testProxyUrl.port;
       redirectedUrl.protocol = testProxyUrl.protocol;
@@ -488,4 +520,16 @@ export class Recorder {
 
     return this.variables[name];
   }
+
+  /**
+   * Set the current skip mode that the recorder should use. All requests made through the recorder
+   * following a call to setSkipMode will be recorded according to the provided skip mode. To
+   * return to recording everything as usual, call setSkipMode("default").
+   *
+   * @param skipMode - the new skip mode to use
+   */
+  setSkipMode(skipMode: SkipMode) {
+    logger.info(`[Recorder#setSkipMode] setting skip mode to ${skipMode}`);
+    this.skipMode = skipMode;
+  }
 }