diff --git a/frontend/jest.config.js b/frontend/jest.config.js
index a31fc6d..310efb5 100644
--- a/frontend/jest.config.js
+++ b/frontend/jest.config.js
@@ -1,6 +1,7 @@
 export default {
   preset: "ts-jest",
   testEnvironment: "node",
+  setupFiles: ["<rootDir>/src/setupTests.ts"],
   transform: {
     "^.+\\.tsx?$": "ts-jest",
   },
diff --git a/frontend/package.json b/frontend/package.json
index 2fb02e3..92bebfc 100644
--- a/frontend/package.json
+++ b/frontend/package.json
@@ -58,6 +58,7 @@
     "@typescript-eslint/parser": "^6.0.0",
     "@vitejs/plugin-react": "^4.0.3",
     "autoprefixer": "^10.4.16",
+    "dotenv": "^16.4.5",
     "eslint": "^8.45.0",
     "eslint-plugin-react-hooks": "^4.6.0",
     "eslint-plugin-react-refresh": "^0.4.3",
diff --git a/frontend/src/setupTests.ts b/frontend/src/setupTests.ts
new file mode 100644
index 0000000..b1ea596
--- /dev/null
+++ b/frontend/src/setupTests.ts
@@ -0,0 +1,3 @@
+// So jest test runner can read env vars from .env file
+import { config } from "dotenv";
+config({ path: ".env.jest" });
diff --git a/frontend/src/tests/qa.test.ts b/frontend/src/tests/qa.test.ts
index 8321543..8fa9972 100644
--- a/frontend/src/tests/qa.test.ts
+++ b/frontend/src/tests/qa.test.ts
@@ -8,7 +8,7 @@ const SCREENSHOTS_PATH = `${REPO_PATH}/qa/results`;
 const IMAGE_PATH = DOWNLOAD_PATH + "/simple_button.png";
 const SCREENSHOT_WITH_IMAGES = `${DOWNLOAD_PATH}/simple_ui_with_image.png`;
 
-describe("Simple Puppeteer Test", () => {
+describe("e2e tests", () => {
   let browser: Browser;
   let page: Page;
 
@@ -105,6 +105,32 @@ describe("Simple Puppeteer Test", () => {
       );
     });
   });
+
+  // Start from code tests - for every model (doesn’t need to be repeated for each stack - fix to HTML Tailwind only)
+  models.forEach((model) => {
+    ["html_tailwind"].forEach((stack) => {
+      it(
+        `Start from code for : ${model}`,
+        async () => {
+          const app = new App(page, stack, model, `update_${model}_${stack}`);
+          await app.init();
+
+          // Generate from screenshot
+          await app.uploadImage(IMAGE_PATH);
+          // Regenerate works for v1
+          await app.regenerate();
+          // Make an update
+          await app.edit("make the header blue", "v2");
+          // Make another update
+          await app.edit("make all text italic", "v3");
+          // Branch off v2 and make an update
+          await app.clickVersion("v2");
+          await app.edit("make all text red", "v4");
+        },
+        90 * 1000
+      );
+    });
+  });
 });
 
 class App {
@@ -132,7 +158,7 @@ class App {
     const setting = {
       openAiApiKey: null,
       openAiBaseURL: null,
-      screenshotOneApiKey: null,
+      screenshotOneApiKey: process.env.TEST_SCREENSHOTONE_API_KEY,
       isImageGenerationEnabled: true,
       editorTheme: "cobalt",
       generatedCodeConfig: this.stack,
@@ -155,21 +181,27 @@ class App {
     });
   }
 
-  async generateFromUrl(url: string) {
-    await this.page.type('input[placeholder="Enter URL"]', url);
+  async _waitUntilVersionIsReady(version: string) {
+    await this.page.waitForNetworkIdle();
+    await this.page.waitForFunction(
+      (version) => document.body.innerText.includes(version),
+      {
+        timeout: 30000,
+      },
+      version
+    );
+    // Wait for 3s so that the HTML and JS has time to render before screenshotting
+    await new Promise((resolve) => setTimeout(resolve, 3000));
+  }
 
+  async generateFromUrl(url: string) {
+    // Type in the URL
+    await this.page.type('input[placeholder="Enter URL"]', url);
     await this._screenshot("typed_url");
 
     // Click the capture button and wait for the code to be generated
     await this.page.click("button.capture-btn");
-    await this.page.waitForNetworkIdle();
-    await this.page.waitForFunction(
-      () => document.body.innerText.includes("v1"),
-      {
-        timeout: 30000,
-      }
-    );
-
+    await this._waitUntilVersionIsReady("v1");
     await this._screenshot("url_result");
   }
 
@@ -179,33 +211,15 @@ class App {
     const fileInput = (await this.page.$(
       ".file-input"
     )) as ElementHandle<HTMLInputElement>;
-
     if (!fileInput) {
       throw new Error("File input element not found");
     }
-
     await fileInput.uploadFile(screenshotPath);
-
-    // Screenshot the first step
-    await this.page.screenshot({
-      path: `${this.screenshotPathPrefix}_image_uploaded.png`,
-    });
+    await this._screenshot("image_uploaded");
 
     // Click the generate button and wait for the code to be generated
-    await this.page.waitForNetworkIdle();
-    await this.page.waitForFunction(
-      () => document.body.innerText.includes("v1"),
-      {
-        timeout: 30000,
-      }
-    );
-
-    // Wait for 1s so that the HTML and JS has time to render before screenshotting
-    await new Promise((resolve) => setTimeout(resolve, 3000));
-
-    await this.page.screenshot({
-      path: `${this.screenshotPathPrefix}_image_results.png`,
-    });
+    await this._waitUntilVersionIsReady("v1");
+    await this._screenshot("image_results");
   }
 
   // Makes a text edit and waits for a new version
@@ -246,16 +260,7 @@ class App {
 
   async regenerate() {
     await this.page.click(".regenerate-btn");
-
-    await this.page.waitForFunction(
-      () => document.body.innerText.includes("v1"),
-      {
-        timeout: 30000,
-      }
-    );
-
-    await this.page.screenshot({
-      path: `${this.screenshotPathPrefix}_regenerate_results.png`,
-    });
+    await this._waitUntilVersionIsReady("v1");
+    await this._screenshot("regenerate_results");
   }
 }
diff --git a/frontend/yarn.lock b/frontend/yarn.lock
index 7a08c1e..7285b73 100644
--- a/frontend/yarn.lock
+++ b/frontend/yarn.lock
@@ -2917,6 +2917,11 @@ dotenv@^16.0.0:
   resolved "https://registry.yarnpkg.com/dotenv/-/dotenv-16.3.1.tgz#369034de7d7e5b120972693352a3bf112172cc3e"
   integrity sha512-IPzF4w4/Rd94bA9imS68tZBaYyBWSCE47V1RGuMrB94iyTOIEwRmVL2x/4An+6mETpLrKJ5hQkB8W4kFAadeIQ==
 
+dotenv@^16.4.5:
+  version "16.4.5"
+  resolved "https://registry.yarnpkg.com/dotenv/-/dotenv-16.4.5.tgz#cdd3b3b604cb327e286b4762e13502f717cb099f"
+  integrity sha512-ZmdL2rui+eB2YwhsWzjInR8LldtZHGDoQ1ugH85ppHKwpUHL7j7rN0Ti9NCnGiQbhaZ11FpR+7ao1dNsmduNUg==
+
 ebml-block@^1.1.2:
   version "1.1.2"
   resolved "https://registry.yarnpkg.com/ebml-block/-/ebml-block-1.1.2.tgz#fd49951b0faf5a3049bdd61c851a76b5e679c290"