clean up test code

frontend/src/App.tsx

@@ -465,7 +465,7 @@ function App() {
                   <div className="flex items-center justify-end gap-x-2 mt-2">
                     <Button
                       onClick={regenerate}
-                      className="flex items-center gap-x-2 dark:text-white dark:bg-gray-700"
+                      className="flex items-center gap-x-2 dark:text-white dark:bg-gray-700 regenerate-btn"
                     >
                       🔄 Regenerate
                     </Button>

frontend/src/tests/qa.test.ts

@@ -5,13 +5,16 @@ import { CodeGenerationModel } from "../lib/models";
 const REPO_PATH = "/Users/abi/Documents/GitHub/screenshot-to-code/frontend";
 const DOWNLOAD_PATH = `${REPO_PATH}/qa`;
 const SCREENSHOTS_PATH = `${REPO_PATH}/qa/results`;
-const IMAGE_PATH = DOWNLOAD_PATH + "/ui_table.png";
+const IMAGE_PATH = DOWNLOAD_PATH + "/simple_button.png";
+const SCREENSHOT_WITH_IMAGES = `${DOWNLOAD_PATH}/simple_ui_with_image.png`;
 
 describe("Simple Puppeteer Test", () => {
   let browser: Browser;
   let page: Page;
 
   const DEBUG = true;
+  const IS_HEADLESS = false;
+
   const stacks = Object.values(Stack).slice(0, DEBUG ? 1 : undefined);
   const models = Object.values(CodeGenerationModel).slice(
     0,
@@ -19,7 +22,7 @@ describe("Simple Puppeteer Test", () => {
   );
 
   beforeAll(async () => {
-    browser = await puppeteer.launch({ headless: false });
+    browser = await puppeteer.launch({ headless: IS_HEADLESS });
     page = await browser.newPage();
     await page.goto("http://localhost:5173/");
 
@@ -40,17 +43,140 @@ describe("Simple Puppeteer Test", () => {
     await browser.close();
   });
 
+  // Create tests
   models.forEach((model) => {
     stacks.forEach((stack) => {
       it(
-        `should load the homepage and check the title for stack: ${stack}`,
+        `Create for : ${model} & ${stack}`,
         async () => {
-          const testId = `${model}_${stack}`;
+          const app = new App(
-          const screenshotPathPrefix = `${SCREENSHOTS_PATH}/${testId}`;
+            page,
-          await setupLocalStorage(page, stack, model);
+            stack,
+            model,
+            `create_screenshot_${model}_${stack}`
+          );
+          await app.init();
+          // Generate from screenshot
+          await app.uploadImage(SCREENSHOT_WITH_IMAGES);
+        },
+        60 * 1000
+      );
 
+      it(
+        `Create from URL for : ${model} & ${stack}`,
+        async () => {
+          const app = new App(
+            page,
+            stack,
+            model,
+            `create_url_${model}_${stack}`
+          );
+          await app.init();
+          // Generate from screenshot
+          await app.generateFromUrl("https://a.picoapps.xyz/design-fear");
+        },
+        60 * 1000
+      );
+    });
+  });
+
+  // Update 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(
+        `Update 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 {
+  private screenshotPathPrefix: string;
+  private page: Page;
+  private stack: string;
+  private model: string;
+
+  // TODOs
+  // - Abstract screenshot functionality
+  // - Abstract waiting for version to be done
+
+  constructor(page: Page, stack: string, model: string, testId: string) {
+    this.page = page;
+    this.stack = stack;
+    this.model = model;
+    this.screenshotPathPrefix = `${SCREENSHOTS_PATH}/${testId}`;
+  }
+
+  async init() {
+    await this.setupLocalStorage();
+  }
+
+  async setupLocalStorage() {
+    const setting = {
+      openAiApiKey: null,
+      openAiBaseURL: null,
+      screenshotOneApiKey: null,
+      isImageGenerationEnabled: true,
+      editorTheme: "cobalt",
+      generatedCodeConfig: this.stack,
+      codeGenerationModel: this.model,
+      isTermOfServiceAccepted: false,
+      accessCode: null,
+    };
+
+    await this.page.evaluate((setting) => {
+      localStorage.setItem("setting", JSON.stringify(setting));
+    }, setting);
+
+    // Reload the page to apply the local storage
+    await this.page.reload();
+  }
+
+  async _screenshot(step: string) {
+    await this.page.screenshot({
+      path: `${this.screenshotPathPrefix}_${step}.png`,
+    });
+  }
+
+  async generateFromUrl(url: string) {
+    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._screenshot("url_result");
+  }
+
+  // Uploads a screenshot and generates the image
+  async uploadImage(screenshotPath: string) {
     // Upload file
-          const fileInput = (await page.$(
+    const fileInput = (await this.page.$(
       ".file-input"
     )) as ElementHandle<HTMLInputElement>;
 
@@ -58,95 +184,44 @@ describe("Simple Puppeteer Test", () => {
       throw new Error("File input element not found");
     }
 
-          await fileInput.uploadFile(IMAGE_PATH);
+    await fileInput.uploadFile(screenshotPath);
 
     // Screenshot the first step
-          await page.screenshot({
+    await this.page.screenshot({
-            path: `${screenshotPathPrefix}_image_uploaded.png`,
+      path: `${this.screenshotPathPrefix}_image_uploaded.png`,
     });
 
     // Click the generate button and wait for the code to be generated
-          await page.waitForNetworkIdle();
+    await this.page.waitForNetworkIdle();
-          await page.waitForFunction(
+    await this.page.waitForFunction(
       () => document.body.innerText.includes("v1"),
       {
         timeout: 30000,
       }
     );
 
-          await page.screenshot({
+    // Wait for 1s so that the HTML and JS has time to render before screenshotting
-            path: `${screenshotPathPrefix}_image_results.png`,
+    await new Promise((resolve) => setTimeout(resolve, 3000));
+
+    await this.page.screenshot({
+      path: `${this.screenshotPathPrefix}_image_results.png`,
     });
-
-          await makeEdit(
-            page,
-            "make the header blue",
-            "v2",
-            screenshotPathPrefix
-          );
-
-          await makeEdit(
-            page,
-            "make all text italic",
-            "v3",
-            screenshotPathPrefix
-          );
-
-          await page.evaluate(() => {
-            document.querySelectorAll("div").forEach((div) => {
-              if (div.innerText.includes("v2")) {
-                div.click();
-              }
-            });
-          });
-
-          await makeEdit(page, "make all text red", "v4", screenshotPathPrefix);
-        },
-        60 * 1000
-      );
-    });
-  });
-});
-
-async function setupLocalStorage(page: Page, stack: string, model: string) {
-  const setting = {
-    openAiApiKey: null,
-    openAiBaseURL: null,
-    screenshotOneApiKey: null,
-    isImageGenerationEnabled: false,
-    editorTheme: "cobalt",
-    generatedCodeConfig: stack,
-    codeGenerationModel: model,
-    isTermOfServiceAccepted: false,
-    accessCode: null,
-  };
-
-  await page.evaluate((setting) => {
-    localStorage.setItem("setting", JSON.stringify(setting));
-  }, setting);
-
-  // Reload the page to apply the local storage
-  await page.reload();
   }
 
-async function makeEdit(
+  // Makes a text edit and waits for a new version
-  page: Page,
+  async edit(edit: string, version: string) {
-  edit: string,
+    await this.page.type(
-  version: string,
-  screenshotPathPrefix: string
-) {
-  await page.type(
       'textarea[placeholder="Tell the AI what to change..."]',
       edit
     );
 
-  await page.screenshot({
+    await this.page.screenshot({
-    path: `${screenshotPathPrefix}_typed_${version}.png`,
+      path: `${this.screenshotPathPrefix}_typed_${version}.png`,
     });
 
-  await page.click(".update-btn");
+    await this.page.click(".update-btn");
 
-  await page.waitForFunction(
+    await this.page.waitForFunction(
       (version: string) => document.body.innerText.includes(version),
       {
         timeout: 30000,
@@ -154,7 +229,33 @@ async function makeEdit(
       version
     );
 
-  await page.screenshot({
+    await this.page.screenshot({
-    path: `${screenshotPathPrefix}_done_${version}.png`,
+      path: `${this.screenshotPathPrefix}_done_${version}.png`,
     });
   }
+
+  async clickVersion(version: string) {
+    await this.page.evaluate((version) => {
+      document.querySelectorAll("div").forEach((div) => {
+        if (div.innerText.includes(version)) {
+          div.click();
+        }
+      });
+    }, version);
+  }
+
+  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`,
+    });
+  }
+}