dotnet/test/common/TakesScreenshotTest.cs

// <copyright file="TakesScreenshotTest.cs" company="Selenium Committers">
// Licensed to the Software Freedom Conservancy (SFC) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The SFC licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.
// </copyright>

using NUnit.Framework;
using OpenQA.Selenium.Environment;
using System;
using System.Collections.Generic;
using System.Drawing;
using System.IO;

namespace OpenQA.Selenium;

[TestFixture]
public class TakesScreenshotTest : DriverTestFixture
{
    [TearDown]
    public void SwitchToTop()
    {
        driver.SwitchTo().DefaultContent();
    }

    [Test]
    public void GetScreenshotAsFile()
    {
        ITakesScreenshot screenshotCapableDriver = driver as ITakesScreenshot;
        if (screenshotCapableDriver == null)
        {
            return;
        }

        driver.Url = simpleTestPage;

        string filename = Path.Combine(Path.GetTempPath(), "snapshot" + new Random().Next().ToString() + ".png");
        Screenshot screenImage = screenshotCapableDriver.GetScreenshot();
        screenImage.SaveAsFile(filename);
        Assert.That(File.Exists(filename), Is.True);
        Assert.That(new FileInfo(filename).Length, Is.GreaterThan(0));
        File.Delete(filename);
    }

    [Test]
    public void GetScreenshotAsBase64()
    {
        ITakesScreenshot screenshotCapableDriver = driver as ITakesScreenshot;
        if (screenshotCapableDriver == null)
        {
            return;
        }

        driver.Url = simpleTestPage;
        Screenshot screenImage = screenshotCapableDriver.GetScreenshot();
        string base64 = screenImage.AsBase64EncodedString;
        Assert.That(base64.Length, Is.GreaterThan(0));
    }

    [Test]
    public void GetScreenshotAsBinary()
    {
        ITakesScreenshot screenshotCapableDriver = driver as ITakesScreenshot;
        if (screenshotCapableDriver == null)
        {
            return;
        }

        driver.Url = simpleTestPage;
        Screenshot screenImage = screenshotCapableDriver.GetScreenshot();
        byte[] bytes = screenImage.AsByteArray;
        Assert.That(bytes.Length, Is.GreaterThan(0));
    }

    [Test]
    [IgnoreBrowser(Browser.IE, "Color comparisons fail on IE")]
    [IgnoreBrowser(Browser.Chrome, "Color comparisons fail on Chrome")]
    [IgnoreBrowser(Browser.Edge, "Color comparisons fail on Edge")]
    [IgnoreBrowser(Browser.Firefox, "Not working properly in RBE, works locally with pinned browsers")]
    public void ShouldCaptureScreenshotOfCurrentViewport()
    {
#if NET8_0
        Assert.Ignore("Skipping test: this framework can not process colors.");
#endif

        ITakesScreenshot screenshotCapableDriver = driver as ITakesScreenshot;
        if (screenshotCapableDriver == null)
        {
            return;
        }

        driver.Url = EnvironmentManager.Instance.UrlBuilder.WhereIs("screen/screen.html");
        Screenshot screenshot = screenshotCapableDriver.GetScreenshot();

        HashSet<string> actualColors = ScanActualColors(screenshot,
                                                    /* stepX in pixels */ 5,
                                                    /* stepY in pixels */ 5);

        HashSet<string> expectedColors = GenerateExpectedColors( /* initial color */ 0x0F0F0F,
                                                             /* color step */ 1000,
                                                             /* grid X size */ 6,
                                                             /* grid Y size */ 6);

        CompareColors(expectedColors, actualColors);
    }

    [Test]
    [IgnoreBrowser(Browser.Chrome, "Color comparisons fail on Chrome")]
    [IgnoreBrowser(Browser.Edge, "Color comparisons fail on Edge")]
    public void ShouldTakeScreenshotsOfAnElement()
    {
#if NET8_0
        Assert.Ignore("Skipping test: this framework can not process colors.");
#endif

        driver.Url = EnvironmentManager.Instance.UrlBuilder.WhereIs("screen/screen.html");
        IWebElement element = driver.FindElement(By.Id("cell11"));

        ITakesScreenshot screenshotCapableElement = element as ITakesScreenshot;
        if (screenshotCapableElement == null)
        {
            return;
        }

        Screenshot screenImage = screenshotCapableElement.GetScreenshot();
        byte[] imageData = screenImage.AsByteArray;
        Assert.That(imageData, Is.Not.Null);
        Assert.That(imageData.Length, Is.GreaterThan(0));

        Color pixelColor = GetPixelColor(screenImage, 1, 1);
        string pixelColorString = FormatColorToHex(pixelColor.ToArgb());
        Assert.That(pixelColorString, Is.EqualTo("#0f12f7"));
    }

    [Test]
    [IgnoreBrowser(Browser.IE, "Color comparisons fail on IE")]
    [IgnoreBrowser(Browser.Chrome, "Color comparisons fail on Chrome")]
    [IgnoreBrowser(Browser.Edge, "Color comparisons fail on Edge")]
    public void ShouldCaptureScreenshotAtFramePage()
    {
#if NET8_0
        Assert.Ignore("Skipping test: this framework can not process colors.");
#endif

        ITakesScreenshot screenshotCapableDriver = driver as ITakesScreenshot;
        if (screenshotCapableDriver == null)
        {
            return;
        }

        driver.Url = EnvironmentManager.Instance.UrlBuilder.WhereIs("screen/screen_frames.html");
        WaitFor(FrameToBeAvailableAndSwitchedTo("frame1"), "Did not switch to frame1");
        WaitFor(ElementToBeVisibleWithId("content"), "Did not find visible element with id content");

        driver.SwitchTo().DefaultContent();
        WaitFor(FrameToBeAvailableAndSwitchedTo("frame2"), "Did not switch to frame2");
        WaitFor(ElementToBeVisibleWithId("content"), "Did not find visible element with id content");

        driver.SwitchTo().DefaultContent();
        WaitFor(TitleToBe("screen test"), "Title was not expected value");
        Screenshot screenshot = screenshotCapableDriver.GetScreenshot();

        HashSet<string> actualColors = ScanActualColors(screenshot,
                                                   /* stepX in pixels */ 5,
                                                   /* stepY in pixels */ 5);

        HashSet<string> expectedColors = GenerateExpectedColors( /* initial color */ 0x0F0F0F,
                                                 /* color step*/ 1000,
                                                 /* grid X size */ 6,
                                                 /* grid Y size */ 6);
        expectedColors.UnionWith(GenerateExpectedColors( /* initial color */ 0xDFDFDF,
                                                 /* color step*/ 1000,
                                                 /* grid X size */ 6,
                                                 /* grid Y size */ 6));

        // expectation is that screenshot at page with frames will be taken for full page
        CompareColors(expectedColors, actualColors);
    }

    [Test]
    [IgnoreBrowser(Browser.IE, "Color comparisons fail on IE")]
    [IgnoreBrowser(Browser.Chrome, "Color comparisons fail on Chrome")]
    [IgnoreBrowser(Browser.Edge, "Color comparisons fail on Edge")]
    public void ShouldCaptureScreenshotAtIFramePage()
    {
#if NET8_0
        Assert.Ignore("Skipping test: this framework can not process colors.");
#endif

        ITakesScreenshot screenshotCapableDriver = driver as ITakesScreenshot;
        if (screenshotCapableDriver == null)
        {
            return;
        }

        driver.Url = EnvironmentManager.Instance.UrlBuilder.WhereIs("screen/screen_iframes.html");

        // Resize the window to avoid scrollbars in screenshot
        Size originalSize = driver.Manage().Window.Size;
        driver.Manage().Window.Size = new Size(1040, 700);

        Screenshot screenshot = screenshotCapableDriver.GetScreenshot();
        driver.Manage().Window.Size = originalSize;

        HashSet<string> actualColors = ScanActualColors(screenshot,
                                                   /* stepX in pixels */ 5,
                                                   /* stepY in pixels */ 5);

        HashSet<string> expectedColors = GenerateExpectedColors( /* initial color */ 0x0F0F0F,
                                                 /* color step*/ 1000,
                                                 /* grid X size */ 6,
                                                 /* grid Y size */ 6);
        expectedColors.UnionWith(GenerateExpectedColors( /* initial color */ 0xDFDFDF,
                                                 /* color step*/ 1000,
                                                 /* grid X size */ 6,
                                                 /* grid Y size */ 6));

        // expectation is that screenshot at page with Iframes will be taken for full page
        CompareColors(expectedColors, actualColors);
    }

    [Test]
    [IgnoreBrowser(Browser.IE, "Color comparisons fail on IE")]
    [IgnoreBrowser(Browser.Firefox, "Color comparisons fail on Firefox")]
    [IgnoreBrowser(Browser.Chrome, "Color comparisons fail on Chrome")]
    [IgnoreBrowser(Browser.Edge, "Color comparisons fail on Edge")]
    public void ShouldCaptureScreenshotAtFramePageAfterSwitching()
    {
#if NET8_0
        Assert.Ignore("Skipping test: this framework can not process colors.");
#endif

        ITakesScreenshot screenshotCapableDriver = driver as ITakesScreenshot;
        if (screenshotCapableDriver == null)
        {
            return;
        }

        driver.Url = EnvironmentManager.Instance.UrlBuilder.WhereIs("screen/screen_frames.html");

        driver.SwitchTo().Frame(driver.FindElement(By.Id("frame2")));

        Screenshot screenshot = screenshotCapableDriver.GetScreenshot();

        HashSet<string> actualColors = ScanActualColors(screenshot,
                                                   /* stepX in pixels */ 5,
                                                   /* stepY in pixels */ 5);

        HashSet<string> expectedColors = GenerateExpectedColors( /* initial color */ 0x0F0F0F,
                                                 /* color step*/ 1000,
                                                 /* grid X size */ 6,
                                                 /* grid Y size */ 6);
        expectedColors.UnionWith(GenerateExpectedColors( /* initial color */ 0xDFDFDF,
                                                 /* color step*/ 1000,
                                                 /* grid X size */ 6,
                                                 /* grid Y size */ 6));

        // expectation is that screenshot at page with frames after switching to a frame
        // will be taken for full page
        CompareColors(expectedColors, actualColors);
    }

    [Test]
    [IgnoreBrowser(Browser.IE, "Color comparisons fail on IE")]
    [IgnoreBrowser(Browser.Firefox, "Color comparisons fail on Firefox")]
    [IgnoreBrowser(Browser.Chrome, "Color comparisons fail on Chrome")]
    [IgnoreBrowser(Browser.Edge, "Color comparisons fail on Edge")]
    public void ShouldCaptureScreenshotAtIFramePageAfterSwitching()
    {
#if NET6_0 || NET8_0
        Assert.Ignore("Skipping test: this framework can not process colors.");
#endif

        ITakesScreenshot screenshotCapableDriver = driver as ITakesScreenshot;
        if (screenshotCapableDriver == null)
        {
            return;
        }

        driver.Url = EnvironmentManager.Instance.UrlBuilder.WhereIs("screen/screen_iframes.html");

        // Resize the window to avoid scrollbars in screenshot
        Size originalSize = driver.Manage().Window.Size;
        driver.Manage().Window.Size = new Size(1040, 700);

        driver.SwitchTo().Frame(driver.FindElement(By.Id("iframe2")));

        Screenshot screenshot = screenshotCapableDriver.GetScreenshot();
        driver.Manage().Window.Size = originalSize;

        HashSet<string> actualColors = ScanActualColors(screenshot,
                                                   /* stepX in pixels */ 5,
                                                   /* stepY in pixels */ 5);

        HashSet<string> expectedColors = GenerateExpectedColors( /* initial color */ 0x0F0F0F,
                                                 /* color step*/ 1000,
                                                 /* grid X size */ 6,
                                                 /* grid Y size */ 6);
        expectedColors.UnionWith(GenerateExpectedColors( /* initial color */ 0xDFDFDF,
                                                 /* color step*/ 1000,
                                                 /* grid X size */ 6,
                                                 /* grid Y size */ 6));

        // expectation is that screenshot at page with Iframes after switching to a Iframe
        // will be taken for full page
        CompareColors(expectedColors, actualColors);
    }

    private string FormatColorToHex(int colorValue)
    {
        string pixelColorString = string.Format("#{0:x2}{1:x2}{2:x2}", (colorValue & 0xFF0000) >> 16, (colorValue & 0x00FF00) >> 8, (colorValue & 0x0000FF));
        return pixelColorString;
    }

    private void CompareColors(HashSet<string> expectedColors, HashSet<string> actualColors)
    {
        // Ignore black and white for further comparison
        actualColors.Remove("#000000");
        actualColors.Remove("#ffffff");

        Assert.That(actualColors, Is.EquivalentTo(expectedColors));
    }

    private HashSet<string> GenerateExpectedColors(int initialColor, int stepColor, int numberOfSamplesX, int numberOfSamplesY)
    {
        HashSet<string> colors = new HashSet<string>();
        int count = 1;
        for (int i = 1; i < numberOfSamplesX; i++)
        {
            for (int j = 1; j < numberOfSamplesY; j++)
            {
                int color = initialColor + (count * stepColor);
                string hex = FormatColorToHex(color);
                colors.Add(hex);
                count++;
            }
        }

        return colors;
    }

    private HashSet<string> ScanActualColors(Screenshot screenshot, int stepX, int stepY)
    {
        HashSet<string> colors = new HashSet<string>();

#if !NET8_0
        try
        {
            Image image = Image.FromStream(new MemoryStream(screenshot.AsByteArray));
            Bitmap bitmap = new Bitmap(image);
            int height = bitmap.Height;
            int width = bitmap.Width;
            Assert.That(width, Is.GreaterThan(0));
            Assert.That(height, Is.GreaterThan(0));

            for (int i = 0; i < width; i = i + stepX)
            {
                for (int j = 0; j < height; j = j + stepY)
                {
                    string hex = FormatColorToHex(bitmap.GetPixel(i, j).ToArgb());
                    colors.Add(hex);
                }
            }
        }
        catch (Exception e)
        {
            Assert.Fail("Unable to get actual colors from screenshot: " + e.Message);
        }

        Assert.That(colors.Count, Is.GreaterThan(0));
#endif

        return colors;
    }

    private Color GetPixelColor(Screenshot screenshot, int x, int y)
    {
        Color pixelColor = Color.Black;

#if !NET8_0
        Image image = Image.FromStream(new MemoryStream(screenshot.AsByteArray));
        Bitmap bitmap = new Bitmap(image);
        pixelColor = bitmap.GetPixel(1, 1);
#endif
        return pixelColor;
    }

    private Func<bool> FrameToBeAvailableAndSwitchedTo(string frameId)
    {
        return () =>
        {
            try
            {
                IWebElement frameElement = driver.FindElement(By.Id(frameId));
                driver.SwitchTo().Frame(frameElement);
            }
            catch (Exception)
            {
                return false;
            }

            return true;
        };
    }

    private Func<bool> ElementToBeVisibleWithId(string elementId)
    {
        return () =>
        {
            try
            {
                IWebElement element = driver.FindElement(By.Id(elementId));
                return element.Displayed;
            }
            catch (Exception)
            {
                return false;
            }
        };
    }

    private Func<bool> TitleToBe(string desiredTitle)
    {
        return () => driver.Title == desiredTitle;
    }
}