Skip to content

Commit

Permalink
Added support with video in selenium
Browse files Browse the repository at this point in the history
  • Loading branch information
JosefShenhav committed Nov 3, 2023
1 parent 928af5a commit a4f85d0
Show file tree
Hide file tree
Showing 4 changed files with 123 additions and 8 deletions.
16 changes: 13 additions & 3 deletions core/testcontainers/core/container.py
Original file line number Diff line number Diff line change
@@ -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__)

Expand All @@ -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
Expand All @@ -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
Expand Down
55 changes: 51 additions & 4 deletions selenium/testcontainers/selenium/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
}


Expand Down Expand Up @@ -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")
Expand All @@ -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)
39 changes: 39 additions & 0 deletions selenium/testcontainers/selenium/video.py
Original file line number Diff line number Diff line change
@@ -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
21 changes: 20 additions & 1 deletion selenium/tests/test_selenium.py
Original file line number Diff line number Diff line change
@@ -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])
Expand All @@ -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"

0 comments on commit a4f85d0

Please sign in to comment.