Rebase and rewrite

Author: JosefShenhav

modules/selenium/testcontainers/selenium/__init__.py

@@ -10,17 +10,20 @@
 #    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 pathlib import Path
 from typing import Optional
 
 import urllib3
+from typing_extensions import Self
 
 from selenium import webdriver
 from selenium.webdriver.common.options import ArgOptions
 from testcontainers.core.container import DockerContainer
+from testcontainers.core.network import Network
 from testcontainers.core.waiting_utils import wait_container_is_ready
+from testcontainers.selenium.video import SeleniumVideoContainer
 
-IMAGES = {"firefox": "selenium/standalone-firefox-debug:latest", "chrome": "selenium/standalone-chrome-debug:latest"}
+IMAGES = {"firefox": "selenium/standalone-firefox:latest", "chrome": "selenium/standalone-chrome:latest"}
 
 
 def get_image_name(capabilities: str) -> str:
@@ -51,6 +54,8 @@ def __init__(
         self.image = image or get_image_name(capabilities)
         self.port = port
         self.vnc_port = vnc_port
+        self.video = None
+        self.__video_network = None
         super().__init__(image=self.image, **kwargs)
         self.with_exposed_ports(self.port, self.vnc_port)
 
@@ -72,3 +77,44 @@ 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, image: Optional[str] = None, video_path: Optional[Path] = None) -> Self:
+        video_path = video_path or Path.cwd()
+
+        self.video = SeleniumVideoContainer(image)
+        video_folder_path = video_path.parent.resolve()
+        self.video.set_videos_host_path(str(video_folder_path))
+
+        if video_path.name:
+            self.video.set_video_name(video_path.name)
+
+        return self
+
+    def start(self) -> 'DockerContainer':
+        if not self.video:
+            super().start()
+            return self
+
+        self.__video_network = Network().__enter__()
+
+        self.with_kwargs(network=self.__video_network.name)
+        super().start()
+
+        self.video \
+            .with_kwargs(network=self.__video_network.name) \
+            .set_selenium_container_host(self.get_wrapped_container().short_id) \
+            .start()
+
+        return self
+
+    def stop(self, force=True, delete_volume=True) -> None:
+        if self.video:
+            # get_wrapped_container().stop -> stop the container
+            # video.stop -> remove the container
+            self.video.get_wrapped_container().stop()
+            self.video.stop(force, delete_volume)
+
+        super().stop(force, delete_volume)
+
+        if self.__video_network:
+            self.__video_network.remove()

modules/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-6.1-20240402"
+
+
+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 set_videos_host_path(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

modules/selenium/tests/test_selenium.py

@@ -1,3 +1,6 @@
+import tempfile
+from pathlib import Path
+
 import pytest
 from selenium.webdriver import DesiredCapabilities
 from selenium.webdriver.common.by import By
@@ -23,3 +26,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.mark.parametrize("caps", [DesiredCapabilities.CHROME, DesiredCapabilities.FIREFOX])
+def test_selenium_video(caps, workdir):
+    video_path = workdir / Path("video.mp4")
+    with BrowserWebDriverContainer(caps).with_video(video_path=video_path) as chrome:
+        chrome.get_driver().get("https://google.com")
+
+    assert video_path.exists(), "Selenium video file does not exists"
+
+
+@pytest.fixture
+def workdir() -> Path:
+    tmpdir = tempfile.TemporaryDirectory()
+    yield Path(tmpdir.name)
+    tmpdir.cleanup()