deflake 1.2.2 → 1.2.4

package/cli.js

@@ -2,6 +2,7 @@
 const yargs = require('yargs/yargs');
 const { hideBin } = require('yargs/helpers');
 const DeFlakeClient = require('./client');
+const Migrator = require('./lib/migrator');
 const { spawn } = require('child_process');
 const fs = require('fs');
 const path = require('path');
@@ -61,9 +62,31 @@ const parser = yargs(hideBin(process.argv))
         description: 'Diagnose your DeFlake installation',
         default: false
     })
+    .option('report', {
+        type: 'boolean',
+        description: 'Automatically open the HTML report after fixing',
+        default: true
+    })
     .command('doctor', 'Diagnose your DeFlake installation', {}, (argv) => {
         runDoctor(argv).then(() => process.exit(0));
     })
+    .command('migrate', 'Migrate Cypress tests to Playwright', {
+        from: { type: 'string', default: 'cypress', description: 'Source framework' },
+        to: { type: 'string', default: 'playwright', description: 'Target framework' },
+        path: { type: 'string', demandOption: true, description: 'Path to Cypress tests' },
+        output: { type: 'string', description: 'Path to output Playwright tests' },
+        ai: { type: 'boolean', default: false, description: 'Enable AI-powered refinement' }
+    }, async (argv) => {
+        try {
+            const client = new DeFlakeClient(argv['api-url']);
+            const migrator = new Migrator({ ...argv, client });
+            await migrator.run();
+            process.exit(0);
+        } catch (err) {
+            console.error(`\x1b[31mMigration failed:\x1b[0m ${err.message}`);
+            process.exit(1);
+        }
+    })
     .version(pkg.version)
     .help();
 
@@ -220,6 +243,28 @@ function parseCypressLogs(fullLog) {
     return specBlocks;
 }
 
+function parsePytestLogs(fullLog) {
+    const cleanLog = stripAnsi(fullLog);
+    const errorBlocks = [];
+    const testHeaderRegex = /_{10,}\s+(.*?)\s+_{10,}/g;
+    let match;
+
+    const indices = [];
+    while ((match = testHeaderRegex.exec(cleanLog)) !== null) {
+        indices.push({ index: match.index, name: match[1] });
+    }
+
+    for (let i = 0; i < indices.length; i++) {
+        const start = indices[i].index;
+        const end = (i + 1 < indices.length) ? indices[i + 1].index : cleanLog.length;
+        errorBlocks.push({
+            name: indices[i].name,
+            content: cleanLog.slice(start, end)
+        });
+    }
+    return errorBlocks;
+}
+
 async function runHealer(logContent, htmlPath, apiUrl, testName, applyFix = false) {
     // Check file size limit (Basic check to avoid 429s on huge files)
     if (htmlPath && fs.existsSync(htmlPath)) {
@@ -321,21 +366,25 @@ function extractFailureLocation(logText) {
 
     const projectName = path.basename(process.cwd());
 
-    // Updated regex to be more flexible with arrows and spaces, and support .cy files
-    const testMatch = logText.match(/^\s*\d+\)\s+\[.*?\]\s+.+?\s+(.*?\.(?:spec|cy)\.(?:ts|js)):(\d+):(\d+)/m);
+    // Updated regex to be more flexible with arrows and spaces, and support .cy/.py files
+    const testMatch = logText.match(/^\s*\d+\)\s+\[.*?\]\s+.+?\s+(.*?\.(?:spec|cy|py)\.(?:ts|js|py)?):(\d+):(\d+)/m) ||
+        logText.match(/^(.*?\.py):(\d+):/m); // Pytest direct format
+
     if (testMatch) {
         loc.specFile = testMatch[1];
         loc.testLine = testMatch[2];
     }
 
-    // Stack Trace Regex - Modified to handle Cypress URLs and webpack paths
-    const stackRegex = /at\s+(?:.*? \()?((?:https?:\/\/.*?\/|webpack:\/\/|[\/~\\]|\.?\.\/|[\w_\-]+\/).*?):(\d+):(\d+)\)?/g;
+    // Stack Trace Regex - Modified to handle Cypress URLs, webpack paths, and Python "File" entries
+    const stackRegex = /at\s+(?:.*? \()?((?:https?:\/\/.*?\/|webpack:\/\/|[\/~\\]|\.?\.\/|[\w_\-]+\/).*?):(\d+):(\d+)\)?|File "(.+?)", line (\d+)/g;
     let match;
     let foundRoot = false;
 
     while ((match = stackRegex.exec(logText)) !== null) {
-        let file = match[1];
-        const line = match[2];
+        let file = match[1] || match[4];
+        const line = match[2] || match[5];
+
+        if (!file) continue;
 
         // Clean Cypress/Browser URLs to just the relative path if possible
         if (file.includes('__cypress/runner')) continue;
@@ -347,16 +396,14 @@ function extractFailureLocation(logText) {
         }
 
         // Strip project name if it's the first segment (common in Cypress/Webpack logs)
-        // e.g. "cypress-poc/cypress/e2e/api/auth.api.cy.ts" -> "cypress/e2e/api/auth.api.cy.ts"
-        // Also handle "./" prefix before project name
         file = file.replace(/^\.\//, '');
         if (file.startsWith(projectName + '/')) {
             file = file.substring(projectName.length + 1);
         }
         file = file.replace(/^\.\//, '');
 
-        if (!foundRoot && !file.includes('node_modules') && !file.includes('cypress_runner')) {
-            loc.rootFile = file.split(/[/\\]/).pop();
+        if (!foundRoot && !file.includes('node_modules') && !file.includes('cypress_runner') && !file.includes('python')) {
+            loc.rootFile = file.split(/[/\\\\]/).pop();
             loc.fullRootPath = path.isAbsolute(file) ? file : path.resolve(process.cwd(), file);
             loc.rootLine = line;
             foundRoot = true;
@@ -368,7 +415,7 @@ function extractFailureLocation(logText) {
     }
 
     // Fallback: If header regex failed but we found a root file that looks like a test
-    if (!loc.specFile && loc.rootFile && (loc.rootFile.includes('.spec.') || loc.rootFile.includes('.test.') || loc.rootFile.includes('.cy.'))) {
+    if (!loc.specFile && loc.rootFile && (loc.rootFile.includes('.spec.') || loc.rootFile.includes('.test.') || loc.rootFile.includes('.cy.') || loc.rootFile.endsWith('.py'))) {
         loc.specFile = loc.rootFile;
         loc.testLine = loc.rootLine;
     }
@@ -532,7 +579,7 @@ async function runDoctor(argv) {
         console.log(`  ✅ Active:    ${C.GREEN}${C.BRIGHT}${capitalized}${C.RESET}`);
     } else {
         console.log(`  ⚠️  ${C.YELLOW}No supported frameworks detected in the current directory.${C.RESET}`);
-        console.log(`     (Checked for Playwright, Cypress, WebdriverIO files)`);
+        console.log(`     (Checked for Playwright, Cypress, WebdriverIO, and Selenium Python files)`);
     }
 
     // 2b. Deep Structure Detection
@@ -707,6 +754,55 @@ async function applySelfHealing(result) {
     }
 }
 
+/**
+ * Automagically opens the framework's native HTML report.
+ */
+function showFrameworkReport() {
+    const framework = DeFlakeClient.detectFramework();
+    console.log(`\n${C.CYAN}📊 Opening HTML Report for ${framework.toUpperCase()}...${C.RESET}`);
+
+    let command = '';
+    const opener = process.platform === 'win32' ? 'start' : 'open';
+
+    if (framework === 'playwright') {
+        command = 'npx playwright show-report';
+    } else if (framework === 'cypress') {
+        // Broad search for common Cypress HTML reports
+        const candidates = [
+            'cypress/reports/html/index.html',
+            'cypress/reports/index.html',
+            'mochawesome-report/mochawesome.html',
+            'reports/index.html'
+        ];
+
+        const found = candidates.find(p => fs.existsSync(p));
+        if (found) {
+            command = `${opener} ${found}`;
+        } else {
+            console.log(`${C.GRAY}ℹ️  Cypress HTML report not found. Verify your reporter configuration (e.g. mochawesome).${C.RESET}`);
+            console.log(`   ${C.GRAY}Looked in: ${candidates.join(', ')}${C.RESET}`);
+            return;
+        }
+    } else if (framework === 'webdriverio') {
+        // WDIO usually uses Allure or spec-reporter
+        if (fs.existsSync('allure-results')) {
+            console.log(`   ${C.GRAY}Allure results detected. Attempting to serve...${C.RESET}`);
+            command = 'npx allure serve allure-results';
+        } else if (fs.existsSync('reports/html/index.html')) {
+            command = `${opener} reports/html/index.html`;
+        } else {
+            console.log(`${C.GRAY}ℹ️  WebdriverIO HTML report not found. If you use Allure, ensure allure-results folder exists.${C.RESET}`);
+            return;
+        }
+    }
+
+    if (command) {
+        // Use inherit if it's a server (like allure serve) or a long-running process
+        const isServer = command.includes('serve') || command.includes('show-report');
+        spawn(command, { shell: true, stdio: isServer ? 'inherit' : 'ignore' });
+    }
+}
+
 async function analyzeFailures(artifacts, fullLog, client) {
     if (artifacts.length === 0) {
         console.log("⚠️  No error artifacts found.");
@@ -733,27 +829,24 @@ async function analyzeFailures(artifacts, fullLog, client) {
         console.log(`   (Processing the first ${limit}...)\n`);
     }
 
-    const framework = artifacts.length > 0 && artifacts[0].htmlPath?.includes('cypress') ? 'Cypress' : 'Playwright';
+    const framework = DeFlakeClient.detectFramework();
     const results = [];
     const processLimit = Math.min(artifacts.length, limit);
     const batchArtifacts = artifacts.slice(0, processLimit);
 
     console.log(`🔍 Analyzing ${batchArtifacts.length} failure(s)...`);
 
-    const playwrightBlocks = (framework === 'Playwright') ? parsePlaywrightLogs(fullLog || "") : [];
-    const cypressBlocks = (framework === 'Cypress') ? parseCypressLogs(fullLog || "") : [];
+    const playwrightBlocks = (framework === 'playwright') ? parsePlaywrightLogs(fullLog || "") : [];
+    const cypressBlocks = (framework === 'cypress') ? parseCypressLogs(fullLog || "") : [];
+    const pytestBlocks = (framework === 'selenium-python') ? parsePytestLogs(fullLog || "") : [];
 
     for (const art of batchArtifacts) {
         let specificLog = fullLog;
 
-        if (framework === 'Cypress' && cypressBlocks.length > 0) {
-            // Match screenshot path to spec log
-            // Screenshot typically: cypress/screenshots/folder/file.cy.ts/Test Name (failed).png
+        if (framework === 'cypress' && cypressBlocks.length > 0) {
             const match = cypressBlocks.find(b => art.htmlPath && art.htmlPath.includes(b.spec));
-            if (match) {
-                specificLog = match.content;
-            }
-        } else if (framework === 'Playwright' && playwrightBlocks.length > 0) {
+            if (match) specificLog = match.content;
+        } else if (framework === 'playwright' && playwrightBlocks.length > 0) {
             // Heuristic matching for Playwright error blocks
             let bestMatch = null;
             let bestScore = -1;
@@ -768,6 +861,10 @@ async function analyzeFailures(artifacts, fullLog, client) {
             if (bestMatch && bestScore > 0) {
                 specificLog = bestMatch.content;
             }
+        } else if (framework === 'selenium-python' && pytestBlocks.length > 0) {
+            // Match pytest block by test name heuristic
+            const match = pytestBlocks.find(b => art.name && b.name.includes(art.name.replace('.png', '')));
+            if (match) specificLog = match.content;
         }
 
         const displayName = (art.name || 'Unknown Artifact').substring(0, 40);
@@ -849,13 +946,18 @@ async function analyzeFailures(artifacts, fullLog, client) {
     } else if (!argv.fix) {
         console.log(`\n${C.BRIGHT}💡 Tip: Use ${C.CYAN}--fix${C.RESET}${C.BRIGHT} to automatically apply these suggested fixes next time.${C.RESET}`);
     }
+
+    // TRIGGER REPORT
+    if (results.length > 0 && argv.fix && argv.report) {
+        showFrameworkReport();
+    }
 }
 
 async function main() {
-    const command = argv._;
+    const command = argv._ || [];
 
-    // If 'doctor' was called, don't proceed to wrapper logic
-    if (command.includes('doctor')) return;
+    // If 'doctor' or 'migrate' was called, don't proceed to wrapper logic
+    if (command.includes('doctor') || command.includes('migrate')) return;
 
     console.log(`🚑 DeFlake JS Client (Batch Mode) - ${C.CYAN}v${pkg.version}${C.RESET}`);
     const client = new DeFlakeClient(argv.apiUrl);

package/client.js

@@ -56,6 +56,7 @@ class DeFlakeClient {
         if (hasFile(/^cypress\.config\.(js|ts|mjs|cjs)$/) || fs.existsSync('cypress.json') || fs.existsSync('cypress')) return 'cypress';
         if (hasFile(/^playwright\.config\.(js|ts|mjs|cjs)$/)) return 'playwright';
         if (hasFile(/^wdio\.conf\.(js|ts|mjs|cjs)$/)) return 'webdriverio';
+        if (fs.existsSync('pytest.ini') || fs.existsSync('conftest.py') || fs.existsSync('tox.ini')) return 'selenium-python';
 
         // 3. Fallback to package.json dependencies
         if (fs.existsSync('package.json')) {
@@ -68,6 +69,14 @@ class DeFlakeClient {
             } catch (e) { }
         }
 
+        // 4. Fallback to Python requirements
+        if (fs.existsSync('requirements.txt')) {
+            try {
+                const reqs = fs.readFileSync('requirements.txt', 'utf8');
+                if (reqs.includes('selenium') || reqs.includes('pytest')) return 'selenium-python';
+            } catch (e) { }
+        }
+
         return 'generic';
     }
 
@@ -147,6 +156,29 @@ class DeFlakeClient {
             return null; // Return null instead of exiting process
         }
     }
+
+    async migrate(sourceCode, framework = 'Playwright') {
+        try {
+            const migrateUrl = this.apiUrl.replace(/\/deflake$/, '/migrate');
+            const payload = {
+                source_code: sourceCode,
+                framework: framework
+            };
+
+            const response = await axios.post(migrateUrl, payload, {
+                headers: {
+                    'Content-Type': 'application/json',
+                    'X-API-KEY': this.apiKey,
+                    'X-Project-Name': this.projectName
+                }
+            });
+
+            return response.data;
+        } catch (error) {
+            console.error(`❌ Migration Refinement Error: ${error.message}`);
+            return { refined_code: sourceCode, status: 'error' };
+        }
+    }
 }
 
 module.exports = DeFlakeClient;

package/lib/mapping.json

@@ -0,0 +1,17 @@
+{
+  "commands": {
+    "visit": "goto",
+    "get": "locator",
+    "type": "fill",
+    "click": "click",
+    "contains": "locator",
+    "url": "url",
+    "wait": "waitForTimeout"
+  },
+  "assertions": {
+    "should": "expect",
+    "be.visible": "toBeVisible",
+    "have.text": "toHaveText",
+    "exist": "toBeVisible"
+  }
+}

package/lib/migrator.js

@@ -0,0 +1,345 @@
+const fs = require('fs');
+const path = require('path');
+const mapping = require('./mapping.json');
+
+// Try to require optional dependencies
+let glob;
+try { glob = require('glob'); } catch (e) { glob = null; }
+let j;
+try { j = require('jscodeshift'); } catch (e) { j = null; }
+
+/**
+ * Migrator handles the conversion of Cypress tests to Playwright 
+ */
+class Migrator {
+    constructor(options = {}) {
+        this.from = options.from || 'cypress';
+        this.to = options.to || 'playwright';
+        this.inputPath = options.path;
+        this.outputPath = options.output || path.join(process.cwd(), 'playwright/tests');
+        this.client = options.client;
+        this.useAI = options.ai || false;
+    }
+
+    async run() {
+        if (!j) {
+            throw new Error("jscodeshift is not installed. Please run 'npm install jscodeshift'");
+        }
+
+        console.log(`\n🚀 Starting migration from ${this.from} to ${this.to}...`);
+        
+        // 1. Scan files
+        const files = await this.scanFiles();
+        console.log(`✔ Found ${files.length} test files`);
+
+        if (files.length === 0) {
+            console.log("⚠️  No files found to migrate.");
+            return;
+        }
+
+        if (!fs.existsSync(this.outputPath)) {
+            fs.mkdirSync(this.outputPath, { recursive: true });
+        }
+
+        // 2. Transform and Generate
+        for (const file of files) {
+            const relativePath = path.relative(this.inputPath, file);
+            const isTS = file.endsWith('.ts') || file.endsWith('.tsx');
+            
+            let targetFilename = path.basename(file);
+            const isTestFile = file.includes('.cy.') || !file.includes('pages/') && !file.includes('support/');
+            
+            if (isTestFile) {
+                targetFilename = targetFilename.replace(/\.(cy\.)?(js|ts)$/, '.spec.ts').replace(/\.(jsx|tsx)$/, '.spec.tsx');
+            } else {
+                targetFilename = targetFilename.replace(/\.js$/, '.ts').replace(/\.jsx$/, '.tsx');
+            }
+            
+            const targetPath = path.join(this.outputPath, path.dirname(relativePath), targetFilename);
+            
+            console.log(`⏳ Converting: ${path.basename(file)}...`);
+            const source = fs.readFileSync(file, 'utf8');
+            
+            // Step A: Deterministic AST Transformation
+            let transformed = this.transform(source, isTS ? 'ts' : 'babylon');
+            
+            // Step B: AI Polishing (if enabled)
+            if (this.useAI && this.client) {
+                console.log(`🧠 Polishing with AI: ${path.basename(file)}...`);
+                const aiResult = await this.client.migrate(transformed, this.to);
+                if (aiResult && aiResult.refined_code) {
+                    transformed = aiResult.refined_code;
+                }
+            }
+            
+            const targetDir = path.dirname(targetPath);
+            if (!fs.existsSync(targetDir)) fs.mkdirSync(targetDir, { recursive: true });
+            
+            fs.writeFileSync(targetPath, transformed);
+        }
+
+
+        console.log(`\n✨ Migration complete!`);
+        console.log(`📂 New tests generated in: ${this.outputPath}`);
+    }
+
+    async scanFiles() {
+        if (glob) {
+            return glob.sync(`${this.inputPath}/**/*.{js,ts,jsx,tsx}`, {
+                ignore: '**/node_modules/**'
+            });
+        }
+        
+        // Fallback simple recursive scan if glob is not available
+        return this.manualScan(this.inputPath);
+    }
+
+    manualScan(dir) {
+        let results = [];
+        const list = fs.readdirSync(dir);
+        list.forEach(file => {
+            const fullPath = path.join(dir, file);
+            const stat = fs.statSync(fullPath);
+            if (stat && stat.isDirectory()) {
+                if (!file.includes('node_modules')) {
+                    results = results.concat(this.manualScan(fullPath));
+                }
+            } else {
+                if (file.endsWith('.js') || file.endsWith('.ts')) {
+                    results.push(fullPath);
+                }
+            }
+        });
+        return results;
+    }
+
+    transform(source, parser = 'babylon') {
+        if (!j) return source;
+        // Use withParser for specific parsers like 'ts'
+        const root = parser === 'ts' ? j.withParser('ts')(source) : j(source);
+
+        // 1. Detect if we are in a Page Object (Class)
+        const classes = root.find(j.ClassDeclaration);
+
+        // 2. Transform Classes for Playwright
+        classes.forEach(path => {
+            const body = path.value.body.body;
+            const hasConstructor = body.some(node => node.type === 'MethodDefinition' && node.kind === 'constructor');
+
+            if (!hasConstructor) {
+                // Inject constructor(page: Page) { this.page = page; }
+                const constructor = j.methodDefinition(
+                    'constructor',
+                    j.identifier('constructor'),
+                    j.functionExpression(
+                        null,
+                        [j.identifier('page')],
+                        j.blockStatement([
+                            j.expressionStatement(
+                                j.assignmentExpression(
+                                    '=',
+                                    j.memberExpression(j.thisExpression(), j.identifier('page')),
+                                    j.identifier('page')
+                                )
+                            )
+                        ])
+                    )
+                );
+                constructor.kind = 'constructor';
+                body.unshift(constructor);
+            }
+        });
+
+        // 3. Command Mapping (Context Aware)
+        // If inside a POM class, use 'this.page', otherwise use 'page'
+        const getContext = (p) => {
+            let parent = p.parentPath;
+            while (parent) {
+                const node = parent.value;
+                if (node.type === 'ClassDeclaration' || node.type === 'ClassMethod' || node.type === 'MethodDefinition') {
+                    return j.memberExpression(j.thisExpression(), j.identifier('page'));
+                }
+                parent = parent.parentPath;
+            }
+            return j.identifier('page');
+        };
+
+
+        // Replace describe -> test.describe, it -> test
+        root.find(j.CallExpression, { callee: { name: 'describe' } })
+            .forEach(p => { 
+                p.get('callee').replace(j.memberExpression(j.identifier('test'), j.identifier('describe'))); 
+            });
+        
+        root.find(j.CallExpression, { callee: { name: 'it' } })
+            .forEach(p => { 
+                p.get('callee').replace(j.identifier('test')); 
+            });
+
+        // Replace cy.get/contains with ctx.locator
+        root.find(j.CallExpression, {
+            callee: {
+                object: { name: 'cy' },
+                property: { name: 'get' }
+            }
+        }).replaceWith(p => {
+            return j.callExpression(
+                j.memberExpression(getContext(p), j.identifier('locator')),
+                p.value.arguments
+            );
+        });
+
+        root.find(j.CallExpression, {
+            callee: {
+                object: { name: 'cy' },
+                property: { name: 'contains' }
+            }
+        }).replaceWith(p => {
+            return j.callExpression(
+                j.memberExpression(getContext(p), j.identifier('locator')),
+                [j.templateLiteral([j.templateElement({ raw: 'text=', cooked: 'text=' }, false)], [p.value.arguments[0]])]
+            );
+        });
+
+        root.find(j.CallExpression, {
+            callee: {
+                object: { name: 'cy' },
+                property: { name: 'visit' }
+            }
+        }).replaceWith(p => {
+            return j.awaitExpression(
+                j.callExpression(
+                    j.memberExpression(getContext(p), j.identifier('goto')),
+                    p.value.arguments
+                )
+            );
+        });
+
+        // Wrap it/test callbacks with async ({ page }) => ...
+        root.find(j.CallExpression, { callee: { name: 'test' } })
+            .forEach(p => {
+                const callback = p.value.arguments[1];
+                if (callback && (callback.type === 'ArrowFunctionExpression' || callback.type === 'FunctionExpression')) {
+                    callback.async = true;
+                    if (callback.params.length === 0) {
+                        callback.params = [j.objectPattern([
+                            j.property.from({
+                                kind: 'init',
+                                key: j.identifier('page'),
+                                value: j.identifier('page'),
+                                shorthand: true
+                            })
+                        ])];
+                    }
+                }
+            });
+
+        // Handle chained commands (click, type, etc)
+        Object.entries(mapping.commands).forEach(([cyCmd, pwCmd]) => {
+            if (['get', 'contains', 'visit'].includes(cyCmd)) return;
+
+            root.find(j.CallExpression, {
+                callee: { property: { name: cyCmd } }
+            }).replaceWith(p => {
+                const call = j.callExpression(
+                    j.memberExpression(p.value.callee.object, j.identifier(pwCmd)),
+                    p.value.arguments
+                );
+                // Wrap with await
+                return j.awaitExpression(call);
+            });
+        });
+
+        // 5. Ensure parent methods are async if they contain await
+        root.find(j.AwaitExpression).forEach(p => {
+            let parent = p.parentPath;
+            while (parent) {
+                const node = parent.value;
+                if (node.type === 'FunctionExpression' || 
+                    node.type === 'ArrowFunctionExpression' || 
+                    node.type === 'ClassMethod' || 
+                    node.type === 'MethodDefinition' ||
+                    node.type === 'ObjectMethod') {
+                    
+                    if (node.type === 'MethodDefinition' || node.type === 'Property') {
+                        if (node.value && (node.value.type === 'FunctionExpression' || node.value.type === 'ArrowFunctionExpression')) {
+                            node.value.async = true;
+                        }
+                    } else {
+                        node.async = true;
+                    }
+                    break;
+                }
+                parent = parent.parentPath;
+            }
+        });
+
+        // Heuristic: If we are calling a method on an object that looks like a POM, await it
+        root.find(j.CallExpression).forEach(p => {
+            if (p.value.callee.type === 'MemberExpression' && 
+                p.value.callee.object.type === 'Identifier' && 
+                p.parentPath.value.type === 'ExpressionStatement') {
+                
+                const objName = p.value.callee.object.name;
+                if (objName.toLowerCase().includes('page') || objName.toLowerCase().includes('login') || objName.toLowerCase().includes('pom')) {
+                    j(p).replaceWith(j.awaitExpression(p.value));
+                }
+            }
+        });
+
+        // 6. Simple TypeScript Type Injections (Basic)
+        if (parser === 'ts') {
+            classes.forEach(path => {
+                const body = path.value.body.body;
+                body.forEach(member => {
+                    if (member.type === 'MethodDefinition' && member.kind === 'constructor') {
+                        const param = member.value.params[0];
+                        if (param && param.name === 'page' && !param.typeAnnotation) {
+                            param.typeAnnotation = j.tsTypeAnnotation(
+                                j.tsTypeReference(j.identifier('Page'))
+                            );
+                        }
+                    }
+                });
+            });
+        }
+
+        // 7. Handle assertions
+        root.find(j.CallExpression, {
+            callee: { property: { name: 'should' } }
+        }).replaceWith(p => {
+            const subject = p.value.callee.object;
+            const assertion = p.value.arguments[0].value;
+            const pwMatcher = mapping.assertions[assertion] || 'toBeVisible';
+
+            // Clean subject: if it's an await expression, get the inner call
+            const cleanSubject = subject.type === 'AwaitExpression' ? subject.argument : subject;
+
+            const call = j.callExpression(
+                j.memberExpression(
+                    j.callExpression(j.identifier('expect'), [cleanSubject]),
+                    j.identifier(pwMatcher)
+                ),
+                p.value.arguments.slice(1)
+            );
+            
+            return j.awaitExpression(call);
+        });
+
+        const result = root.toSource();
+        // Add imports
+        if (!result.includes('@playwright/test')) {
+            const imports = parser === 'ts' ? '{ test, expect, Page }' : '{ test, expect }';
+            return `import ${imports} from '@playwright/test';\n\n${result}`;
+        } else if (parser === 'ts' && result.includes('{ test, expect }') && !result.includes('Page')) {
+            return result.replace('{ test, expect }', '{ test, expect, Page }');
+        }
+
+        return result;
+
+    }
+}
+
+module.exports = Migrator;
+
+

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "deflake",
-  "version": "1.2.2",
+  "version": "1.2.4",
   "description": "AI-powered self-healing tool for Playwright, Cypress, and WebdriverIO tests.",
   "main": "client.js",
   "bin": {

package/sample-cypress/login.cy.js

@@ -0,0 +1,9 @@
+describe('Login', () => {
+  it('should login', () => {
+    cy.visit('/login')
+    cy.get('#email').type('test@test.com')
+    cy.get('#password').type('123456')
+    cy.get('button').click()
+    cy.contains('Dashboard').should('be.visible')
+  })
+})

package/sample-cypress-pom/login_pom.cy.js

@@ -0,0 +1,11 @@
+const LoginPage = require('../pages/LoginPage')
+
+describe('Login POM', () => {
+  const loginPage = new LoginPage()
+
+  it('should login with POM', () => {
+    loginPage.visit()
+    loginPage.login('test@test.com', '123456')
+    cy.contains('Dashboard').should('be.visible')
+  })
+})

package/sample-cypress-pom/pages/LoginPage.js

@@ -0,0 +1,12 @@
+class LoginPage {
+  visit() {
+    cy.visit('/login')
+  }
+
+  login(email, password) {
+    cy.get('#email').type(email)
+    cy.get('#password').type(password)
+    cy.get('button').click()
+  }
+}
+module.exports = LoginPage;

package/sample-cypress-ts/login_ts.cy.js

@@ -0,0 +1,11 @@
+import LoginPage from './pages/LoginPage'
+
+describe('Login TS', () => {
+  const loginPage = new LoginPage()
+
+  it('should login with TS', () => {
+    loginPage.visit()
+    loginPage.login('test@test.com', '123456')
+    cy.contains('Dashboard').should('be.visible')
+  })
+})

package/sample-cypress-ts/pages/LoginPage.ts

@@ -0,0 +1,12 @@
+class LoginPage {
+  visit() {
+    cy.visit('/login')
+  }
+
+  login(email: string, password: string) {
+    cy.get('#email').type(email)
+    cy.get('#password').type(password)
+    cy.get('button').click()
+  }
+}
+export default LoginPage;

package/sample-playwright/login.spec.spec.ts

@@ -0,0 +1,15 @@
+import { test, expect } from '@playwright/test';
+
+test.describe('Login', () => {
+  test('should login', async (
+    {
+      page
+    }
+  ) => {
+    await page.goto('/login')
+    (await page.locator('#email')).type('test@test.com')
+    (await page.locator('#password')).type('123456')
+    (await page.locator('button')).click()
+    await expect(page.locator('Dashboard')).toBeVisible()
+  })
+})

package/sample-playwright-ai-v2/login_pom.spec.ts

@@ -0,0 +1,17 @@
+import { test, expect } from '@playwright/test';
+
+const LoginPage = require('../pages/LoginPage')
+
+test.describe('Login POM', () => {
+  const loginPage = new LoginPage()
+
+  test('should login with POM', async (
+    {
+      page
+    }
+  ) => {
+    await loginPage.visit()
+    await loginPage.login('test@test.com', '123456')
+    await expect(page.locator(`text=${'Dashboard'}`)).toBeVisible()
+  })
+})

package/sample-playwright-ai-v2/pages/LoginPage.ts

@@ -0,0 +1,18 @@
+import { test, expect } from '@playwright/test';
+
+class LoginPage {
+  constructor(page) {
+    this.page = page;
+  }
+
+  async visit() {
+    await this.page.goto('/login')
+  }
+
+  async login(email, password) {
+    await this.page.locator('#email').fill(email)
+    await this.page.locator('#password').fill(password)
+    await this.page.locator('button').click()
+  }
+}
+module.exports = LoginPage;

package/sample-playwright-final/login.spec.ts

@@ -0,0 +1,15 @@
+import { test, expect } from '@playwright/test';
+
+test.describe('Login', () => {
+  test('should login', async (
+    {
+      page
+    }
+  ) => {
+    await page.goto('/login')
+    await page.locator('#email').fill('test@test.com')
+    await page.locator('#password').fill('123456')
+    await page.locator('button').click()
+    await expect(page.locator(`text=${'Dashboard'}`)).toBeVisible()
+  })
+})

package/sample-playwright-pom/login_pom.spec.ts

@@ -0,0 +1,17 @@
+import { test, expect } from '@playwright/test';
+
+const LoginPage = require('../pages/LoginPage')
+
+test.describe('Login POM', () => {
+  const loginPage = new LoginPage()
+
+  test('should login with POM', async (
+    {
+      page
+    }
+  ) => {
+    loginPage.visit()
+    loginPage.login('test@test.com', '123456')
+    await expect(page.locator(`text=${'Dashboard'}`)).toBeVisible()
+  })
+})

package/sample-playwright-pom/pages/LoginPage.spec.ts

@@ -0,0 +1,18 @@
+import { test, expect } from '@playwright/test';
+
+class LoginPage {
+  constructor(page) {
+    this.page = page;
+  }
+
+  visit() {
+    await this.page.goto('/login')
+  }
+
+  login(email, password) {
+    await this.page.locator('#email').fill(email)
+    await this.page.locator('#password').fill(password)
+    await this.page.locator('button').click()
+  }
+}
+module.exports = LoginPage;

package/sample-playwright-pom-v2/login_pom.spec.ts

@@ -0,0 +1,17 @@
+import { test, expect } from '@playwright/test';
+
+const LoginPage = require('../pages/LoginPage')
+
+test.describe('Login POM', () => {
+  const loginPage = new LoginPage()
+
+  test('should login with POM', async (
+    {
+      page
+    }
+  ) => {
+    await loginPage.visit()
+    await loginPage.login('test@test.com', '123456')
+    await expect(page.locator(`text=${'Dashboard'}`)).toBeVisible()
+  })
+})

package/sample-playwright-pom-v2/pages/LoginPage.spec.ts

@@ -0,0 +1,18 @@
+import { test, expect } from '@playwright/test';
+
+class LoginPage {
+  constructor(page) {
+    this.page = page;
+  }
+
+  async visit() {
+    await this.page.goto('/login')
+  }
+
+  login(email, password) {
+    await this.page.locator('#email').fill(email)
+    await this.page.locator('#password').fill(password)
+    await this.page.locator('button').click()
+  }
+}
+module.exports = LoginPage;

package/sample-playwright-pom-v3/login_pom.spec.ts

@@ -0,0 +1,17 @@
+import { test, expect } from '@playwright/test';
+
+const LoginPage = require('../pages/LoginPage')
+
+test.describe('Login POM', () => {
+  const loginPage = new LoginPage()
+
+  test('should login with POM', async (
+    {
+      page
+    }
+  ) => {
+    await loginPage.visit()
+    await loginPage.login('test@test.com', '123456')
+    await expect(page.locator(`text=${'Dashboard'}`)).toBeVisible()
+  })
+})

package/sample-playwright-pom-v3/pages/LoginPage.spec.ts

@@ -0,0 +1,18 @@
+import { test, expect } from '@playwright/test';
+
+class LoginPage {
+  constructor(page) {
+    this.page = page;
+  }
+
+  async visit() {
+    await this.page.goto('/login')
+  }
+
+  login(email, password) {
+    await this.page.locator('#email').fill(email)
+    await this.page.locator('#password').fill(password)
+    await this.page.locator('button').click()
+  }
+}
+module.exports = LoginPage;

package/sample-playwright-pom-v4/login_pom.spec.ts

@@ -0,0 +1,17 @@
+import { test, expect } from '@playwright/test';
+
+const LoginPage = require('../pages/LoginPage')
+
+test.describe('Login POM', () => {
+  const loginPage = new LoginPage()
+
+  test('should login with POM', async (
+    {
+      page
+    }
+  ) => {
+    await loginPage.visit()
+    await loginPage.login('test@test.com', '123456')
+    await expect(page.locator(`text=${'Dashboard'}`)).toBeVisible()
+  })
+})

package/sample-playwright-pom-v4/pages/LoginPage.spec.ts

@@ -0,0 +1,18 @@
+import { test, expect } from '@playwright/test';
+
+class LoginPage {
+  constructor(page) {
+    this.page = page;
+  }
+
+  async visit() {
+    await this.page.goto('/login')
+  }
+
+  async login(email, password) {
+    await this.page.locator('#email').fill(email)
+    await this.page.locator('#password').fill(password)
+    await this.page.locator('button').click()
+  }
+}
+module.exports = LoginPage;

package/sample-playwright-ts-v2/login_ts.spec.ts

@@ -0,0 +1,17 @@
+import { test, expect } from '@playwright/test';
+
+import LoginPage from './pages/LoginPage'
+
+test.describe('Login TS', () => {
+  const loginPage = new LoginPage()
+
+  test('should login with TS', async (
+    {
+      page
+    }
+  ) => {
+    await loginPage.visit()
+    await loginPage.login('test@test.com', '123456')
+    await expect(page.locator(`text=${'Dashboard'}`)).toBeVisible()
+  })
+})

package/sample-playwright-ts-v2/pages/LoginPage.ts

@@ -0,0 +1,18 @@
+import { test, expect, Page } from '@playwright/test';
+
+class LoginPage {
+  constructor(page: Page) {
+    this.page = page;
+  }
+
+  async visit() {
+    await this.page.goto('/login')
+  }
+
+  async login(email: string, password: string) {
+    await this.page.locator('#email').fill(email)
+    await this.page.locator('#password').fill(password)
+    await this.page.locator('button').click()
+  }
+}
+export default LoginPage;

package/sample-playwright-v2/login.spec.ts

@@ -0,0 +1,11 @@
+import { test, expect } from '@playwright/test';
+
+describe('Login', () => {
+  it('should login', () => {
+    await page.goto('/login')
+    await page.locator('#email').fill('test@test.com')
+    await page.locator('#password').fill('123456')
+    await page.locator('button').click()
+    await expect(page.locator(`text=${'Dashboard'}`)).toBeVisible()
+  })
+})