Added support with video in selenium

Author: JosefShenhav

core/testcontainers/core/container.py

@@ -1,11 +1,12 @@
-from docker.models.containers import Container
 import os
 from typing import Iterable, Optional, Tuple
 
-from .waiting_utils import wait_container_is_ready
+from docker.models.containers import Container
+
 from .docker_client import DockerClient
 from .exceptions import ContainerStartException
 from .utils import setup_logger, inside_container, is_arm
+from .waiting_utils import wait_container_is_ready
 
 logger = setup_logger(__name__)
 
@@ -28,11 +29,16 @@ def __init__(self, image: str, docker_client_kw: Optional[dict] = None, **kwargs
         self.volumes = {}
         self.image = image
         self._docker = DockerClient(**(docker_client_kw or {}))
+        self.network_name = None
         self._container = None
         self._command = None
         self._name = None
         self._kwargs = kwargs
 
+    @property
+    def name(self):
+        return self._container.name
+
     def with_env(self, key: str, value: str) -> 'DockerContainer':
         self.env[key] = value
         return self
@@ -55,12 +61,16 @@ def maybe_emulate_amd64(self) -> 'DockerContainer':
             return self.with_kwargs(platform='linux/amd64')
         return self
 
+    def set_network_name(self, network_name: str) -> 'DockerContainer':
+        self.network_name = network_name
+        return self
+
     def start(self) -> 'DockerContainer':
         logger.info("Pulling image %s", self.image)
         docker_client = self.get_docker_client()
         self._container = docker_client.run(
             self.image, command=self._command, detach=True, environment=self.env, ports=self.ports,
-            name=self._name, volumes=self.volumes, **self._kwargs
+            name=self._name, volumes=self.volumes, network=self.network_name, **self._kwargs
         )
         logger.info("Container started: %s", self._container.short_id)
         return self

selenium/testcontainers/selenium/__init__.py

@@ -10,17 +10,22 @@
 #    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 #    License for the specific language governing permissions and limitations
 #    under the License.
+import uuid
+from pathlib import Path
+from typing import Optional
 
+import urllib3
 from selenium import webdriver
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_container_is_ready
-from typing import Optional
-import urllib3
 
+from .video import SeleniumVideoContainer
+
+EMPTY_PATH = "."
 
 IMAGES = {
-    "firefox": "selenium/standalone-firefox-debug:latest",
-    "chrome": "selenium/standalone-chrome-debug:latest"
+    "firefox": "selenium/standalone-firefox:4.13.0-20231004",
+    "chrome": "selenium/standalone-chrome:4.13.0-20231004"
 }
 
 
@@ -53,6 +58,8 @@ def __init__(self, capabilities: str, image: Optional[str] = None, port: int = 4
         self.vnc_port = vnc_port
         super(BrowserWebDriverContainer, self).__init__(image=self.image, **kwargs)
         self.with_exposed_ports(self.port, self.vnc_port)
+        self.video = None
+        self.network = None
 
     def _configure(self) -> None:
         self.with_env("no_proxy", "localhost")
@@ -71,3 +78,43 @@ def get_connection_url(self) -> str:
         ip = self.get_container_host_ip()
         port = self.get_exposed_port(self.port)
         return f'http://{ip}:{port}/wd/hub'
+
+    def with_video(self, video_path: Path = Path.cwd()) -> 'DockerContainer':
+        self.video = SeleniumVideoContainer()
+
+        target_video_path = video_path.parent
+        if target_video_path.samefile(EMPTY_PATH):
+            target_video_path = Path.cwd()
+        self.video.save_videos(str(target_video_path))
+
+        if video_path.name:
+            self.video.set_video_name(video_path.name)
+
+        return self
+
+    def start(self) -> 'DockerContainer':
+        if self.video:
+            self.network = self._docker.client.networks.create(str(uuid.uuid1()))
+            self.set_network_name(self.network.name)
+            self.video.set_network_name(self.network.name)
+
+            super().start()
+
+            self.video.set_selenium_container_host(self.get_wrapped_container().short_id)
+            self.video.start()
+
+            return self
+
+        super().start()
+        return self
+
+    def stop(self, force=True, delete_volume=True) -> None:
+        if self.video:
+            # Video need to stop before remove
+            self.video.get_wrapped_container().stop()
+            self.video.stop(force, delete_volume)
+
+        super().stop(force, delete_volume)
+
+        if self.network:
+            self.get_docker_client().client.api.remove_network(self.network.id)

selenium/testcontainers/selenium/video.py

@@ -0,0 +1,39 @@
+#
+#    Licensed 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.
+from typing import Optional
+
+from testcontainers.core.container import DockerContainer
+
+VIDEO_DEFAULT_IMAGE = "selenium/video:ffmpeg-4.3.1-20231004"
+
+
+class SeleniumVideoContainer(DockerContainer):
+    """
+    Selenium video container.
+    """
+
+    def __init__(self, image: Optional[str] = None, **kwargs) -> None:
+        self.image = image or VIDEO_DEFAULT_IMAGE
+        super().__init__(image=self.image, **kwargs)
+
+    def set_video_name(self, video_name: str) -> 'DockerContainer':
+        self.with_env("FILE_NAME", video_name)
+        return self
+
+    def save_videos(self, host_path: str) -> 'DockerContainer':
+        self.with_volume_mapping(host_path, "/videos", "rw")
+        return self
+
+    def set_selenium_container_host(self, host: str) -> 'DockerContainer':
+        self.with_env("DISPLAY_CONTAINER_NAME", host)
+        return self

selenium/tests/test_selenium.py

@@ -1,7 +1,10 @@
+import tempfile
+from pathlib import Path
+
 import pytest
 from selenium.webdriver import DesiredCapabilities
-from testcontainers.selenium import BrowserWebDriverContainer
 from testcontainers.core.utils import is_arm
+from testcontainers.selenium import BrowserWebDriverContainer
 
 
 @pytest.mark.parametrize("caps", [DesiredCapabilities.CHROME, DesiredCapabilities.FIREFOX])
@@ -20,3 +23,19 @@ def test_selenium_custom_image():
     chrome = BrowserWebDriverContainer(DesiredCapabilities.CHROME, image=image)
     assert "image" in dir(chrome), "`image` attribute was not instantialized."
     assert chrome.image == image, "`image` attribute was not set to the user provided value"
+
+
+@pytest.fixture
+def workdir() -> Path:
+    tmpdir = tempfile.TemporaryDirectory()
+    yield Path(tmpdir.name)
+    tmpdir.cleanup()
+
+
+@pytest.mark.parametrize("caps", [DesiredCapabilities.CHROME])
+def test_selenium_video(caps, workdir):
+    video_path = workdir / Path("video.mp4")
+    with BrowserWebDriverContainer(caps).with_video(video_path) as chrome:
+        chrome.get_driver()
+
+    assert video_path.exists(), "Selenium video file does not exists"