fix: Add selenium video support #6

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 if video_path.suffix else video_path
+ self.video.set_videos_host_path(str(video_folder_path.resolve()))
+
+ 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,7 @@
+import os
+import tempfile
+from pathlib import Path
+
 import pytest
 from selenium.webdriver import DesiredCapabilities
 from selenium.webdriver.common.by import By
@@ -23,3 +27,21 @@ 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 exist"
+
+
+@pytest.fixture
+def workdir() -> Path:
+ tmpdir = tempfile.TemporaryDirectory()
+ # Enable write permissions for the Docker user container.
+ os.chmod(tmpdir.name, 0o777)
+ yield Path(tmpdir.name)
+ tmpdir.cleanup()