package snappi.loader;

import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.IOException;

import javax.imageio.ImageIO;

import org.openqa.selenium.OutputType;
import org.openqa.selenium.SearchContext;
import org.openqa.selenium.TakesScreenshot;

import snappi.image.WebImage;


public class SeleniumLoader 
{
    public static WebImage load(TakesScreenshot driverOrElement) {
      byte[] bytes = driverOrElement.getScreenshotAs(OutputType.BYTES);
      
      try {
        BufferedImage image = ImageIO.read(new ByteArrayInputStream(bytes));
        
        WebImage webImage = new WebImage(image);
        if (driverOrElement instanceof SearchContext) {
          webImage.setSearchContext((SearchContext) driverOrElement);
        }
        
        return webImage;
      }
      catch(IOException e) {
        throw new IllegalArgumentException("An error occured while reading the image.");
      }
    }
}

@@ -275,16 +275,55 @@
 eenshot;
+%0Aimport org.openqa.selenium.WebElement;
 %0A%0Aimport
@@ -640,16 +640,185 @@
 ebImage 
+webImage;%0A        if(driverOrElement instanceof WebElement) %7B%0A          webImage = new WebImage(image, (WebElement) driverOrElement);%0A        %7D%0A        else %7B%0A          
 webImage
@@ -841,16 +841,35 @@
 image);%0A
+        %7D%0A        %0A
         

Set the relative element