11
11
# License for the specific language governing permissions and limitations
12
12
# under the License.
13
13
from pathlib import Path
14
- from typing import Optional
14
+ from typing import Any , Optional
15
15
16
16
import urllib3
17
17
from typing_extensions import Self
26
26
IMAGES = {"firefox" : "selenium/standalone-firefox:latest" , "chrome" : "selenium/standalone-chrome:latest" }
27
27
28
28
29
- def get_image_name (capabilities : str ) -> str :
29
+ def get_image_name (capabilities : dict [ str , Any ] ) -> str :
30
30
return IMAGES [capabilities ["browserName" ]]
31
31
32
32
@@ -48,9 +48,16 @@ class BrowserWebDriverContainer(DockerContainer):
48
48
"""
49
49
50
50
def __init__ (
51
- self , capabilities : str , image : Optional [str ] = None , port : int = 4444 , vnc_port : int = 5900 , ** kwargs
51
+ self ,
52
+ capabilities : dict [str , Any ],
53
+ options : Optional [ArgOptions ] = None ,
54
+ image : Optional [str ] = None ,
55
+ port : int = 4444 ,
56
+ vnc_port : int = 5900 ,
57
+ ** kwargs ,
52
58
) -> None :
53
59
self .capabilities = capabilities
60
+ self .options = options
54
61
self .image = image or get_image_name (capabilities )
55
62
self .port = port
56
63
self .vnc_port = vnc_port
@@ -65,7 +72,7 @@ def _configure(self) -> None:
65
72
66
73
@wait_container_is_ready (urllib3 .exceptions .HTTPError )
67
74
def _connect (self ) -> webdriver .Remote :
68
- options = ArgOptions ()
75
+ options = ArgOptions () if self . options is None else self . options
69
76
for key , value in self .capabilities .items ():
70
77
options .set_capability (key , value )
71
78
return webdriver .Remote (command_executor = (self .get_connection_url ()), options = options )
@@ -78,6 +85,10 @@ def get_connection_url(self) -> str:
78
85
port = self .get_exposed_port (self .port )
79
86
return f"http://{ ip } :{ port } /wd/hub"
80
87
88
+ def with_options (self , options : Optional [ArgOptions ]):
89
+ self .options = options
90
+ return self
91
+
81
92
def with_video (self , image : Optional [str ] = None , video_path : Optional [Path ] = None ) -> Self :
82
93
video_path = video_path or Path .cwd ()
83
94
0 commit comments