How to Build a Custom Visual Testing Tool with Selenium: A Practical Guide

from selenium import webdriver
import os

def capture_screenshot(url, filename):
    """Capture screenshot of a webpage"""
    driver = webdriver.Chrome()
    driver.get(url)
    
    # Ensure screenshot directory exists
    os.makedirs("screenshots", exist_ok=True)
    screenshot_path = f"screenshots/{filename}"
    
    driver.save_screenshot(screenshot_path)
    driver.quit()
    return screenshot_path

from PIL import Image, ImageChops
import math

def compare_images(baseline_path, current_path, diff_path=None, threshold=0.95):
    """
    Compare two images with similarity threshold
    Returns: (is_similar, similarity_score)
    """
    baseline = Image.open(baseline_path).convert('RGB')
    current = Image.open(current_path).convert('RGB')
    
    # Check dimensions
    if baseline.size != current.size:
        return False, 0
    
    # Calculate difference
    diff = ImageChops.difference(baseline, current)
    diff_pixels = sum(
        sum(1 for pixel in diff.getdata() if any(c > 0 for c in pixel))
    )
    total_pixels = baseline.size[0] * baseline.size[1]
    similarity = 1 - (diff_pixels / total_pixels)
    
    # Save diff image if needed
    if diff_path and diff_pixels > 0:
        diff.save(diff_path)
    
    return similarity >= threshold, similarity

import json
from datetime import datetime

class BaselineManager:
    def __init__(self, baseline_dir="baselines"):
        self.baseline_dir = baseline_dir
        os.makedirs(baseline_dir, exist_ok=True)
        
    def save_baseline(self, name, image_path):
        """Save a new baseline with metadata"""
        timestamp = datetime.now().isoformat()
        baseline_path = f"{self.baseline_dir}/{name}.png"
        metadata = {
            "created": timestamp,
            "source": image_path
        }
        
        # Save image
        Image.open(image_path).save(baseline_path)
        
        # Save metadata
        with open(f"{baseline_dir}/{name}.json", 'w') as f:
            json.dump(metadata, f)
            
        return baseline_path

def compare_regions(baseline_path, current_path, regions, threshold=0.95):
    """
    Compare specific regions of images
    regions: List of (x, y, width, height) tuples
    """
    baseline = Image.open(baseline_path)
    current = Image.open(current_path)
    results = []
    
    for region in regions:
        x, y, w, h = region
        baseline_crop = baseline.crop((x, y, x+w, y+h))
        current_crop = current.crop((x, y, x+w, y+h))
        
        is_similar, score = compare_images(
            baseline_crop, current_crop, threshold=threshold
        )
        results.append((region, is_similar, score))
    
    return results

def mask_dynamic_regions(image_path, regions, output_path=None):
    """
    Mask dynamic content regions with black rectangles
    """
    img = Image.open(image_path)
    draw = ImageDraw.Draw(img)
    
    for region in regions:
        x, y, w, h = region
        draw.rectangle((x, y, x+w, y+h), fill='black')
    
    if output_path:
        img.save(output_path)
    return img

def run_visual_test(url, test_name, threshold=0.95):
    """Complete visual test workflow"""
    # Setup
    bm = BaselineManager()
    current_path = capture_screenshot(url, f"current_{test_name}.png")
    
    # Check if baseline exists
    baseline_path = f"baselines/{test_name}.png"
    if not os.path.exists(baseline_path):
        print(f"Creating new baseline for {test_name}")
        bm.save_baseline(test_name, current_path)
        return True
    
    # Compare images
    diff_path = f"diffs/diff_{test_name}.png"
    is_similar, score = compare_images(
        baseline_path, current_path, diff_path, threshold
    )
    
    # Generate report
    report = {
        "test_name": test_name,
        "passed": is_similar,
        "similarity_score": score,
        "diff_image": diff_path if not is_similar else None,
        "timestamp": datetime.now().isoformat()
    }
    
    return report

import unittest

class VisualTestCase(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.bm = BaselineManager()
        
    def test_homepage_layout(self):
        report = run_visual_test(
            "https://example.com", 
            "homepage_desktop",
            threshold=0.98
        )
        self.assertTrue(report['passed'], 
            f"Visual regression detected. Similarity: {report['similarity_score']}")

def generate_html_report(test_reports):
    """Generate visual test HTML report"""
    html = """
    <html><head><title>Visual Test Report</title></head>
    <body><h1>Visual Test Results</h1>
    <table border="1">
    <tr>
        <th>Test</th>
        <th>Status</th>
        <th>Similarity</th>
        <th>Diff</th>
    </tr>
    """
    
    for report in test_reports:
        status = "PASS" if report['passed'] else "FAIL"
        color = "green" if report['passed'] else "red"
        diff_link = f'<a href="{report["diff_image"]}">View</a>' if report["diff_image"] else "None"
        
        html += f"""
        <tr>
            <td>{report['test_name']}</td>
            <td style="color:{color}">{status}</td>
            <td>{report['similarity_score']:.2%}</td>
            <td>{diff_link}</td>
        </tr>
        """
    
    html += "</table></body></html>"
    return html