feat: Support tunneling

src/api/llm/ws.ts

@@ -38,6 +38,8 @@ export function generateTest(
   assertions: string[],
   prevGoal: string = '',
   creds: Credentials,
+  tunnelName: string,
+  tunnelOwner: string,
 ): [
   WebSocket,
   Observable<
@@ -81,6 +83,8 @@ export function generateTest(
       platform: platform,
       platform_version: platformVersion,
       region: region,
+      tunnel_name: tunnelName,
+      tunnel_owner: tunnelOwner,
     };
 
     ws.onopen = () => {
@@ -90,6 +94,8 @@ export function generateTest(
           data: {
             username: username,
             access_key: accessKey,
+            tunnel_name: tunnelName,
+            tunnel_owner: tunnelOwner,
             app_name: appName,
             goal: goal,
             num_steps: maxTestSteps,

src/editors/TestGenerationPanel.ts

@@ -243,14 +243,14 @@
 
   /**
    * Sets up an event listener to listen for messages passed from the webview context and
-   * executes code based on the message that is recieved.
+   * executes code based on the message that is received.
    *
    * @param webview A reference to the extension webview
    * @param context A reference to the extension context
    */
   private _setWebviewMessageListener(webview: Webview) {
     webview.onDidReceiveMessage(
       async (message: any) => {
         switch (message.action) {
           case 'ready':
             this._msgQueue.ready();
@@ -265,6 +265,8 @@
                 message.data.platform,
                 message.data.platform_version,
                 message.data.assertions,
+                message.data.tunnel_name,
+                message.data.tunnel_owner,
                 '',
               );
             } catch (e) {
@@ -447,6 +449,8 @@
     platform: Platform,
     platformVersion: string,
     assertions: string[],
+    tunnel_name: string,
+    tunnel_owner: string,
     prevGoal: string,
   ) {
     const creds = this.getCredentials();
@@ -512,6 +516,8 @@
       assertions,
       prevGoal,
       creds,
+      tunnel_name,
+      tunnel_owner,
     );
     this._socket = ws;
 
@@ -558,7 +564,7 @@
 }
 
 class MessageQueue {
   private _messages: any[];
   private _ready: boolean;
   private _webview: Webview;
 
@@ -568,7 +574,7 @@
     this._webview = webview;
   }
 
   enqueue(msg: any) {
     if (!this._ready) {
       this._messages.push(msg);
       return;

src/panels/test-generation.ts

@@ -147,7 +147,7 @@
    */
   private subscribeToWebviewEvents(webview: vscode.Webview) {
     webview.onDidReceiveMessage(
       async (message: any) => {
         switch (message.action) {
           case 'stop-generation':
             this.stopTestGeneration();
@@ -162,6 +162,8 @@
                 message.data.platform,
                 message.data.platform_version,
                 message.data.assertions,
+                message.data.tunnel_name,
+                message.data.tunnel_owner,
                 '',
               );
             } catch (e) {
@@ -308,6 +310,8 @@
     platform: Platform,
     platformVersion: string,
     assertions: string[],
+    tunnelName: string,
+    tunnelOwner: string,
     prevGoal: string,
   ) {
     const creds = this.getCredentials();
@@ -367,6 +371,8 @@
       assertions,
       prevGoal,
       creds,
+      tunnelName,
+      tunnelOwner,
     );
     this.socket = ws;
 

src/types.ts

@@ -22,6 +22,8 @@ export interface TestRecord {
   screen_width?: number;
   screen_height?: number;
   scriptiq_llm_version?: string;
+  tunnel_name?: string;
+  tunnel_owner?: string;
 }
 
 export interface TestStep {

webview-ui/test-generation/src/App.tsx

@@ -26,7 +26,7 @@
   const [showAdditionalSettings, setShowAdditionalSettings] = useState(false);
 
   useEffect(() => {
     function handler(event: any) {
       const message = event.data as PostedMessage; // The json data that the extension sent
 
       switch (message.action) {
@@ -106,8 +106,16 @@
     });
   }, []);
 
-  const { appName, assertions, testGoal, maxSteps, platform, status, steps } =
-    state;
+  const {
+    appName,
+    assertions,
+    testGoal,
+    maxSteps,
+    platform,
+    status,
+    steps,
+    tunnel,
+  } = state;
 
   return (
     <main>
@@ -166,6 +174,34 @@
 
         {showAdditionalSettings && (
           <section className="inputs">
+            <VSCodeTextField
+              placeholder="Sauce Connect tunnel name"
+              onInput={(e) => {
+                if (e.target && 'value' in e.target) {
+                  dispatch({
+                    type: 'setTunnelName',
+                    value: e.target.value as string,
+                  });
+                }
+              }}
+              value={state.tunnel?.name ?? ''}
+            >
+              Tunnel Name
+            </VSCodeTextField>
+            <VSCodeTextField
+              placeholder="Sauce Connect tunnel owner"
+              onInput={(e) => {
+                if (e.target && 'value' in e.target) {
+                  dispatch({
+                    type: 'setTunnelOwner',
+                    value: e.target.value as string,
+                  });
+                }
+              }}
+              value={state.tunnel?.owner ?? ''}
+            >
+              Tunnel Owner
+            </VSCodeTextField>
             <VSCodeTextField
               value={maxSteps.toString()}
               onInput={(e) => {
@@ -260,6 +296,8 @@
                   devices: state.device ? [state.device] : [],
                   platform: platform.name,
                   platform_version: platform.version,
+                  tunnel_name: tunnel?.name,
+                  tunnel_owner: tunnel?.owner,
                 },
               });
             }}
@@ -377,6 +415,8 @@
                           state.job?.platform.version ?? '',
                           state.region ?? '',
                           state.platform.name,
+                          state.tunnel?.name ?? '',
+                          state.tunnel?.owner ?? '',
                           steps,
                         );
 

webview-ui/test-generation/src/codegen/base.ts

@@ -19,7 +19,7 @@
    * @returns the line of code
    */
   abstract genCodeLine(
     bestIdentifier: any,
     action: string,
     opts?: {
       // No idea what this is
@@ -54,6 +54,8 @@
     platform_version: string,
     region: string,
     platform: string,
+    tunnel_name: string,
+    tunnel_owner: string,
   ): string;
 
   /**
@@ -82,7 +84,9 @@
     platform_version: string,
     region: string,
     platform: string,
+    tunnel_name: string,
+    tunnel_owner: string,
     steps: any[],
   ): string;
 
   abstract noOptionComment(): string;

webview-ui/test-generation/src/codegen/java.ts

@@ -75,6 +75,8 @@ export class AppiumJava extends AbstractBaseGenerator {
     platform_version: string,
     region: string,
     platform: string,
+    tunnel_name: string,
+    tunnel_owner: string,
   ) {
     const automationName = platform == 'Android' ? 'UiAutomator2' : 'xcuitest';
     const split_goal = this.splitComments(goal, false, `Goal: `);
@@ -133,6 +135,8 @@ public class TestGenerationAssistantTest {
         sauceOptions.setCapability("name", name.getMethodName() + ": " + "${goal}");
         sauceOptions.setCapability("username", System.getenv("SAUCE_USERNAME"));
         sauceOptions.setCapability("accessKey", System.getenv("SAUCE_ACCESS_KEY"));
+        sauceOptions.setCapability("tunnelName", "${tunnel_name}");
+        sauceOptions.setCapability("tunnelOwner", "${tunnel_owner}");
 
         capabilities.setCapability("sauce:options", sauceOptions);
     }
@@ -190,6 +194,8 @@ public class TestGenerationAssistantTest {
     platform_version: string,
     region: string,
     platform: string,
+    tunnel_name: string,
+    tunnel_owner: string,
     steps: any[],
   ) {
     const headerText = this.scriptHeaderCode(
@@ -199,6 +205,8 @@ public class TestGenerationAssistantTest {
       platform_version,
       region,
       platform,
+      tunnel_name,
+      tunnel_owner,
     );
 
     let codeStepText = '';

webview-ui/test-generation/src/codegen/python.ts

@@ -74,6 +74,8 @@ export class AppiumPython extends AbstractBaseGenerator {
     platform_version: string,
     region: string,
     platform: string,
+    tunnel_name: string,
+    tunnel_owner: string,
   ) {
     const automationName = platform == 'Android' ? 'UiAutomator2' : 'xcuitest';
     return `import os
@@ -100,6 +102,8 @@ options = UiAutomator2Options().load_capabilities(
             "username": os.environ["SAUCE_USERNAME"],
             "accessKey": os.environ["SAUCE_ACCESS_KEY"],
             "appiumVersion": "latest",
+            "tunnelName": "${tunnel_name}",
+            "tunnelOwner": "${tunnel_owner}",
         },
     },
 )
@@ -152,6 +156,8 @@ driver.quit()\n`;
     platform_version: string,
     region: string,
     platform: string,
+    tunnel_name: string,
+    tunnel_owner: string,
     steps: any[],
   ) {
     const headerText = this.scriptHeaderCode(
@@ -161,6 +167,8 @@ driver.quit()\n`;
       platform_version,
       region,
       platform,
+      tunnel_name,
+      tunnel_owner,
     );
 
     let codeStepText = '';

webview-ui/test-generation/src/state.ts

@@ -47,6 +47,11 @@ export interface Step {
   screen_descs: string[];
 }
 
+export interface Tunnel {
+  name?: string;
+  owner?: string;
+}
+
 export interface State {
   appName: string;
   testGoal: string;
@@ -71,6 +76,8 @@ export interface State {
   credentials?: Credentials;
 
   language: 'python' | 'java';
+
+  tunnel: Tunnel;
 }
 
 export const initialState: State = {
@@ -91,6 +98,10 @@ export const initialState: State = {
   generationState: 'idle',
   status: '',
   language: 'python',
+  tunnel: {
+    name: '',
+    owner: '',
+  },
 };
 
 export type Action =
@@ -103,6 +114,8 @@ export type Action =
   | { type: 'setStatus'; value: State['status'] }
   | { type: 'setGenerationState'; value: State['generationState'] }
   | { type: 'setDevice'; value: string }
+  | { type: 'setTunnelName'; value: State['tunnel']['name'] }
+  | { type: 'setTunnelOwner'; value: State['tunnel']['owner'] }
   | { type: 'showTestRecord'; value: { testRecord: TestRecord; votes: Vote[] } }
   | { type: 'loadNewRecord'; value: TestRecord }
   | { type: 'addAssertion'; value: { key: string } }
@@ -341,6 +354,10 @@ export const reducer = (current: State, action: Action): State => {
           key: uuidv4(),
           value,
         })),
+        tunnel: {
+          name: testRecord.tunnel_name,
+          owner: testRecord.tunnel_owner,
+        },
       };
     }
     case 'showTestRecord': {
@@ -403,6 +420,10 @@ export const reducer = (current: State, action: Action): State => {
           key: uuidv4(),
           value,
         })),
+        tunnel: {
+          name: testRecord.tunnel_name,
+          owner: testRecord.tunnel_owner,
+        },
       };
     }
     case 'showVideo': {
@@ -413,6 +434,24 @@ export const reducer = (current: State, action: Action): State => {
         credentials: action.value.credentials,
       };
     }
+    case 'setTunnelName': {
+      return {
+        ...current,
+        tunnel: {
+          ...current.tunnel,
+          name: action.value,
+        },
+      };
+    }
+    case 'setTunnelOwner': {
+      return {
+        ...current,
+        tunnel: {
+          ...current.tunnel,
+          owner: action.value,
+        },
+      };
+    }
     default:
       return current;
   }