Add support for a SeleniumBase Context Manager

Author: mdmintz

seleniumbase/__init__.py

@@ -7,8 +7,11 @@
 from seleniumbase.core.browser_launcher import get_driver  # noqa
 from seleniumbase.fixtures import js_utils  # noqa
 from seleniumbase.fixtures import page_actions  # noqa
+from seleniumbase.fixtures import page_utils  # noqa
 from seleniumbase.fixtures.base_case import BaseCase  # noqa
 from seleniumbase.masterqa.master_qa import MasterQA  # noqa
+from seleniumbase.plugins.sb_manager import SB  # noqa
+from seleniumbase.plugins.driver_manager import Driver  # noqa
 
 if sys.version_info[0] >= 3:
     from seleniumbase import translate  # noqa

seleniumbase/behave/behave_sb.py

@@ -180,6 +180,7 @@ def get_configured_sb(context):
     sb.visual_baseline = False
     sb.window_size = None
     sb.maximize_option = False
+    sb.is_context_manager = False
     sb.save_screenshot_after_test = False
     sb.timeout_multiplier = None
     sb.pytest_html_report = None
@@ -215,6 +216,7 @@ def get_configured_sb(context):
     sb.proxy_pac_url = None
     sb.swiftshader = False
     sb.ad_block_on = False
+    sb.is_nosetest = False
     sb.highlights = None
     sb.interval = None
     sb.cap_file = None
@@ -826,6 +828,10 @@ def get_configured_sb(context):
     sb_config.headless = sb.headless
     sb_config.headless_active = False
     sb_config.headed = sb.headed
+    sb_config.is_behave = True
+    sb_config.is_pytest = False
+    sb_config.is_nosetest = False
+    sb_config.is_context_manager = False
     sb_config.window_size = sb.window_size
     sb_config.maximize_option = sb.maximize_option
     sb_config.xvfb = sb.xvfb

seleniumbase/core/browser_launcher.py

@@ -876,7 +876,7 @@ def validate_proxy_string(proxy_string):
 
 
 def get_driver(
-    browser_name,
+    browser_name=None,
     headless=False,
     locale_code=None,
     use_grid=False,
@@ -921,7 +921,19 @@ def get_driver(
     device_width=None,
     device_height=None,
     device_pixel_ratio=None,
+    browser=None,  # A duplicate of browser_name to avoid confusion
 ):
+    if not browser_name:
+        if browser:
+            browser_name = browser
+        else:
+            browser_name = "chrome"  # The default if not specified
+    browser_name = browser_name.lower()
+    if headless2 and browser_name == constants.Browser.FIREFOX:
+        headless2 = False  # Only for Chromium
+        headless = True
+    if uc_subprocess and not undetectable:
+        undetectable = True
     proxy_auth = False
     proxy_user = None
     proxy_pass = None

seleniumbase/core/log_helper.py

@@ -330,6 +330,24 @@ def get_test_id(test):
             scenario_name = scenario_name.split(" -- @")[0]
         test_id = "%s:%s => %s" % (file_name, line_num, scenario_name)
         return test_id
+    elif hasattr(test, "is_context_manager") and test.is_context_manager:
+        filename = test.__class__.__module__.split(".")[-1] + ".py"
+        classname = test.__class__.__name__
+        methodname = test._testMethodName
+        context_id = None
+        if filename == "base_case.py" or methodname == "runTest":
+            import traceback
+
+            stack_base = traceback.format_stack()[0].split(", in ")[0]
+            test_base = stack_base.split(", in ")[0].split(os.sep)[-1]
+            if hasattr(test, "cm_filename") and test.cm_filename:
+                filename = test.cm_filename
+            else:
+                filename = test_base.split('"')[0]
+            classname = "SB"
+            methodname = ".py:" + test_base.split(", line ")[-1]
+            context_id = filename.split(".")[0] + methodname + ":" + classname
+            return context_id
     test_id = None
     try:
         test_id = get_test_name(test)

seleniumbase/fixtures/base_case.py

@@ -3433,6 +3433,7 @@ def get_new_driver(
             device_width=d_width,
             device_height=d_height,
             device_pixel_ratio=d_p_r,
+            browser=browser_name,
         )
         self._drivers_list.append(new_driver)
         self._drivers_browser_map[new_driver] = browser_name
@@ -4571,6 +4572,23 @@ def __process_recorded_actions(self):
         filename = self.__get_filename()
         classname = self.__class__.__name__
         methodname = self._testMethodName
+        context_filename = None
+        if (
+            hasattr(sb_config, "is_context_manager")
+            and sb_config.is_context_manager
+            and (filename == "base_case.py" or methodname == "runTest")
+        ):
+            import traceback
+
+            stack_base = traceback.format_stack()[0].split(os.sep)[-1]
+            test_base = stack_base.split(", in ")[0]
+            if hasattr(self, "cm_filename") and self.cm_filename:
+                filename = self.cm_filename
+            else:
+                filename = test_base.split('"')[0]
+            classname = "SB"
+            methodname = "test_line_" + test_base.split(", line ")[-1]
+            context_filename = filename.split(".")[0] + "_rec.py"
         if hasattr(self, "is_behave") and self.is_behave:
             classname = sb_config.behave_feature.name
             classname = classname.replace("/", " ").replace(" & ", " ")
@@ -4627,6 +4645,8 @@ def __process_recorded_actions(self):
             file_name = sb_config.behave_scenario.filename.replace(".", "_")
             file_name = file_name.split("/")[-1].split("\\")[-1]
             file_name = file_name + "_rec.py"
+        elif context_filename:
+            file_name = context_filename
         file_path = os.path.join(recordings_folder, file_name)
         out_file = codecs.open(file_path, "w+", "utf-8")
         out_file.writelines("\r\n".join(data))
@@ -13245,6 +13265,11 @@ def __has_exception(self):
         has_exception = False
         if hasattr(sys, "last_traceback") and sys.last_traceback is not None:
             has_exception = True
+        elif hasattr(self, "is_context_manager") and self.is_context_manager:
+            if self.with_testing_base and self._has_failure:
+                return True
+            else:
+                return False
         elif python3 and hasattr(self, "_outcome"):
             if hasattr(self._outcome, "errors") and self._outcome.errors:
                 has_exception = True
@@ -13277,6 +13302,22 @@ def __get_test_id(self):
             scenario_name = scenario_name.replace(" ", "_")
             test_id = "%s.%s" % (file_name, scenario_name)
             return test_id
+        elif hasattr(self, "is_context_manager") and self.is_context_manager:
+            filename = self.__class__.__module__.split(".")[-1] + ".py"
+            methodname = self._testMethodName
+            context_id = None
+            if filename == "base_case.py" or methodname == "runTest":
+                import traceback
+
+                stack_base = traceback.format_stack()[0].split(", in ")[0]
+                test_base = stack_base.split(", in ")[0].split(os.sep)[-1]
+                if hasattr(self, "cm_filename") and self.cm_filename:
+                    filename = self.cm_filename
+                else:
+                    filename = test_base.split('"')[0]
+                methodname = ".line_" + test_base.split(", line ")[-1]
+                context_id = filename.split(".")[0] + methodname
+                return context_id
         test_id = "%s.%s.%s" % (
             self.__class__.__module__,
             self.__class__.__name__,

seleniumbase/plugins/base_plugin.py

@@ -7,6 +7,7 @@
 import sys
 import time
 from nose.plugins import Plugin
+from seleniumbase import config as sb_config
 from seleniumbase.config import settings
 from seleniumbase.core import download_helper
 from seleniumbase.core import log_helper
@@ -200,10 +201,12 @@ def configure(self, options, conf):
         archive_logs = options.archive_logs
         log_helper.log_folder_setup(log_path, archive_logs)
         download_helper.reset_downloads_folder()
+        sb_config.is_nosetest = True
         if self.report_on:
             report_helper.clear_out_old_report_logs(archive_past_runs=False)
 
     def beforeTest(self, test):
+        sb_config._context_of_runner = False  # Context Manager Compatibility
         variables = self.options.variables
         if variables and type(variables) is str and len(variables) > 0:
             bad_input = False

seleniumbase/plugins/pytest_plugin.py

@@ -1285,6 +1285,9 @@ def pytest_configure(config):
     sb_config.item_count_skipped = 0
     sb_config.item_count_untested = 0
     sb_config.is_pytest = True
+    sb_config.is_behave = False
+    sb_config.is_nosetest = False
+    sb_config.is_context_manager = False
     sb_config.pytest_config = config
     sb_config.browser = config.getoption("browser")
     if sb_config._browser_shortcut:
@@ -1302,7 +1305,10 @@ def pytest_configure(config):
     sb_config.device_metrics = config.getoption("device_metrics")
     sb_config.headless = config.getoption("headless")
     sb_config.headless2 = config.getoption("headless2")
-    if sb_config.browser not in ["chrome", "edge"]:
+    if sb_config.headless2 and sb_config.browser == "firefox":
+        sb_config.headless2 = False  # Only for Chromium browsers
+        sb_config.headless = True  # Firefox has regular headless
+    elif sb_config.browser not in ["chrome", "edge"]:
         sb_config.headless2 = False  # Only for Chromium browsers
     sb_config.headed = config.getoption("headed")
     sb_config.xvfb = config.getoption("xvfb")
@@ -1624,6 +1630,7 @@ def pytest_collection_finish(session):
     """This runs after item collection is finalized.
     https://docs.pytest.org/en/stable/reference.html
     """
+    sb_config._context_of_runner = False  # Context Manager Compatibility
     if "--co" in sys_argv or "--collect-only" in sys_argv:
         return
     if len(session.items) > 0 and not sb_config._multithreaded:
@@ -1713,14 +1720,20 @@ def pytest_runtest_teardown(item):
             if hasattr(self, "xvfb") and self.xvfb:
                 if self.headless_active and "--pdb" not in sys_argv:
                     if hasattr(self, "display") and self.display:
+                        self.headless_active = False
+                        sb_config.headless_active = False
                         self.display.stop()
             elif hasattr(self, "headless") and self.headless:
                 if self.headless_active and "--pdb" not in sys_argv:
                     if hasattr(self, "display") and self.display:
+                        self.headless_active = False
+                        sb_config.headless_active = False
                         self.display.stop()
             elif hasattr(self, "headless2") and self.headless2:
                 if self.headless_active and "--pdb" not in sys_argv:
                     if hasattr(self, "display") and self.display:
+                        self.headless_active = False
+                        sb_config.headless_active = False
                         self.display.stop()
         except Exception:
             pass