1111#    License for the specific language governing permissions and limitations 
1212#    under the License. 
1313from  pathlib  import  Path 
14- from  typing  import  Optional 
14+ from  typing  import  Any ,  Optional 
1515
1616import  urllib3 
1717from  typing_extensions  import  Self 
2626IMAGES  =  {"firefox" : "selenium/standalone-firefox:latest" , "chrome" : "selenium/standalone-chrome:latest" }
2727
2828
29- def  get_image_name (capabilities : str ) ->  str :
29+ def  get_image_name (capabilities : dict [ str ,  Any ] ) ->  str :
3030    return  IMAGES [capabilities ["browserName" ]]
3131
3232
@@ -48,9 +48,16 @@ class BrowserWebDriverContainer(DockerContainer):
4848    """ 
4949
5050    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 ,
5258    ) ->  None :
5359        self .capabilities  =  capabilities 
60+         self .options  =  options 
5461        self .image  =  image  or  get_image_name (capabilities )
5562        self .port  =  port 
5663        self .vnc_port  =  vnc_port 
@@ -65,7 +72,7 @@ def _configure(self) -> None:
6572
6673    @wait_container_is_ready (urllib3 .exceptions .HTTPError ) 
6774    def  _connect (self ) ->  webdriver .Remote :
68-         options  =  ArgOptions ()
75+         options  =  ArgOptions ()  if   self . options   is   None   else   self . options 
6976        for  key , value  in  self .capabilities .items ():
7077            options .set_capability (key , value )
7178        return  webdriver .Remote (command_executor = (self .get_connection_url ()), options = options )
@@ -78,6 +85,10 @@ def get_connection_url(self) -> str:
7885        port  =  self .get_exposed_port (self .port )
7986        return  f"http://{ ip } { port }  
8087
88+     def  with_options (self , options : Optional [ArgOptions ]):
89+         self .options  =  options 
90+         return  self 
91+ 
8192    def  with_video (self , image : Optional [str ] =  None , video_path : Optional [Path ] =  None ) ->  Self :
8293        video_path  =  video_path  or  Path .cwd ()
8394
0 commit comments