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,121 +43,219 @@ 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}`;
-          await setupLocalStorage(page, stack, model);
-
-          // Upload file
-          const fileInput = (await page.$(
-            ".file-input"
-          )) as ElementHandle<HTMLInputElement>;
-
-          if (!fileInput) {
-            throw new Error("File input element not found");
-          }
-
-          await fileInput.uploadFile(IMAGE_PATH);
-
-          // Screenshot the first step
-          await page.screenshot({
-            path: `${screenshotPathPrefix}_image_uploaded.png`,
-          });
-
-          // Click the generate button and wait for the code to be generated
-          await page.waitForNetworkIdle();
-          await page.waitForFunction(
-            () => document.body.innerText.includes("v1"),
-            {
-              timeout: 30000,
-            }
-          );
-
-          await page.screenshot({
-            path: `${screenshotPathPrefix}_image_results.png`,
-          });
-
-          await makeEdit(
             page,
-            "make the header blue",
+            stack,
-            "v2",
+            model,
-            screenshotPathPrefix
+            `create_screenshot_${model}_${stack}`
           );
+          await app.init();
+          // Generate from screenshot
+          await app.uploadImage(SCREENSHOT_WITH_IMAGES);
+        },
+        60 * 1000
+      );
 
-          await makeEdit(
+      it(
+        `Create from URL for : ${model} & ${stack}`,
+        async () => {
+          const app = new App(
             page,
-            "make all text italic",
+            stack,
-            "v3",
+            model,
-            screenshotPathPrefix
+            `create_url_${model}_${stack}`
           );
-
+          await app.init();
-          await page.evaluate(() => {
+          // Generate from screenshot
-            document.querySelectorAll("div").forEach((div) => {
+          await app.generateFromUrl("https://a.picoapps.xyz/design-fear");
-              if (div.innerText.includes("v2")) {
-                div.click();
-              }
-            });
-          });
-
-          await makeEdit(page, "make all text red", "v4", screenshotPathPrefix);
         },
         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
+      );
+    });
+  });
 });
 
-async function setupLocalStorage(page: Page, stack: string, model: string) {
+class App {
-  const setting = {
+  private screenshotPathPrefix: string;
-    openAiApiKey: null,
+  private page: Page;
-    openAiBaseURL: null,
+  private stack: string;
-    screenshotOneApiKey: null,
+  private model: string;
-    isImageGenerationEnabled: false,
-    editorTheme: "cobalt",
-    generatedCodeConfig: stack,
-    codeGenerationModel: model,
-    isTermOfServiceAccepted: false,
-    accessCode: null,
-  };
 
-  await page.evaluate((setting) => {
+  // TODOs
-    localStorage.setItem("setting", JSON.stringify(setting));
+  // - Abstract screenshot functionality
-  }, setting);
+  // - Abstract waiting for version to be done
 
-  // Reload the page to apply the local storage
+  constructor(page: Page, stack: string, model: string, testId: string) {
-  await page.reload();
+    this.page = page;
-}
+    this.stack = stack;
-
+    this.model = model;
-async function makeEdit(
+    this.screenshotPathPrefix = `${SCREENSHOTS_PATH}/${testId}`;
-  page: Page,
+  }
-  edit: string,
+
-  version: string,
+  async init() {
-  screenshotPathPrefix: string
+    await this.setupLocalStorage();
-) {
+  }
-  await page.type(
+
-    'textarea[placeholder="Tell the AI what to change..."]',
+  async setupLocalStorage() {
-    edit
+    const setting = {
-  );
+      openAiApiKey: null,
-
+      openAiBaseURL: null,
-  await page.screenshot({
+      screenshotOneApiKey: null,
-    path: `${screenshotPathPrefix}_typed_${version}.png`,
+      isImageGenerationEnabled: true,
-  });
+      editorTheme: "cobalt",
-
+      generatedCodeConfig: this.stack,
-  await page.click(".update-btn");
+      codeGenerationModel: this.model,
-
+      isTermOfServiceAccepted: false,
-  await page.waitForFunction(
+      accessCode: null,
-    (version: string) => document.body.innerText.includes(version),
+    };
-    {
+
-      timeout: 30000,
+    await this.page.evaluate((setting) => {
-    },
+      localStorage.setItem("setting", JSON.stringify(setting));
-    version
+    }, setting);
-  );
+
-
+    // Reload the page to apply the local storage
-  await page.screenshot({
+    await this.page.reload();
-    path: `${screenshotPathPrefix}_done_${version}.png`,
+  }
-  });
+
+  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 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`,
+    });
+
+    // 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`,
+    });
+  }
+
+  // Makes a text edit and waits for a new version
+  async edit(edit: string, version: string) {
+    await this.page.type(
+      'textarea[placeholder="Tell the AI what to change..."]',
+      edit
+    );
+
+    await this.page.screenshot({
+      path: `${this.screenshotPathPrefix}_typed_${version}.png`,
+    });
+
+    await this.page.click(".update-btn");
+
+    await this.page.waitForFunction(
+      (version: string) => document.body.innerText.includes(version),
+      {
+        timeout: 30000,
+      },
+      version
+    );
+
+    await this.page.screenshot({
+      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`,
+    });
+  }
 }