Refactor optional features

Author: mdmintz

seleniumbase/core/mysql.py

@@ -1,23 +1,30 @@
-"""
-Wrapper for MySQL DB functions to make life easier.
-"""
-
-import sys
-import time
-from seleniumbase import config as sb_config
-from seleniumbase.config import settings
-from seleniumbase.core import settings_parser
+"""Wrapper for MySQL DB functions."""
 
 
 class DatabaseManager:
-    """
-    This class wraps MySQL database methods for easy use.
-    """
+    """This class wraps MySQL database methods for easy use."""
 
     def __init__(self, database_env="test", conf_creds=None):
-        """
-        Create a connection to the MySQL DB.
-        """
+        """Create a connection to the MySQL DB."""
+        import fasteners
+        import sys
+        import time
+        from importlib.util import find_spec
+        from seleniumbase import config as sb_config
+        from seleniumbase.config import settings
+        from seleniumbase.core import settings_parser
+        from seleniumbase.fixtures import constants
+        from seleniumbase.fixtures import shared_utils
+
+        pip_find_lock = fasteners.InterProcessLock(
+            constants.PipInstall.FINDLOCK
+        )
+        with pip_find_lock:  # Prevent multi-processes mode issues
+            if not find_spec("pymysql"):
+                if sys.version_info >= (3, 6):
+                    shared_utils.pip_install("pymysql", version="1.0.2")
+                else:
+                    shared_utils.pip_install("pymysql", version="0.10.1")
         import pymysql
 
         db_server = settings.DB_HOST
@@ -70,27 +77,21 @@ def __init__(self, database_env="test", conf_creds=None):
             raise Exception("Unable to connect to Database after 3 retries.")
 
     def query_fetch_all(self, query, values):
-        """
-        Executes a db query, gets all the values, and closes the connection.
-        """
+        """Execute db query, get all the values, and close the connection."""
         self.cursor.execute(query, values)
         retval = self.cursor.fetchall()
         self.__close_db()
         return retval
 
     def query_fetch_one(self, query, values):
-        """
-        Executes a db query, gets the first value, and closes the connection.
-        """
+        """Execute db query, get the first value, and close the connection."""
         self.cursor.execute(query, values)
         retval = self.cursor.fetchone()
         self.__close_db()
         return retval
 
     def execute_query(self, query, values):
-        """
-        Executes a query to the test_db and closes the connection afterwards.
-        """
+        """Execute db query, close the connection, and return the results."""
         retval = self.cursor.execute(query, values)
         self.__close_db()
         return retval

seleniumbase/core/s3_manager.py

@@ -1,17 +1,12 @@
-"""
-Methods for uploading/managing files on Amazon S3.
-"""
-from seleniumbase.config import settings
-from seleniumbase.fixtures import shared_utils
+"""Methods for uploading/managing files on Amazon S3."""
 
 already_uploaded_files = []
 
 
 class S3LoggingBucket(object):
-    """
-    A class for uploading log files from tests to Amazon S3.
-    Those files can then be shared easily.
-    """
+    """A class for uploading log files from tests to Amazon S3.
+    Those files can then be shared easily."""
+    from seleniumbase.config import settings
 
     def __init__(
         self,
@@ -20,23 +15,26 @@ def __init__(
         selenium_access_key=settings.S3_SELENIUM_ACCESS_KEY,
         selenium_secret_key=settings.S3_SELENIUM_SECRET_KEY,
     ):
-        try:
-            from boto.s3.connection import S3Connection
-        except Exception:
-            shared_utils.pip_install("boto", version="2.49.0")
-            from boto.s3.connection import S3Connection
+        import fasteners
+        from importlib.util import find_spec
+        from seleniumbase.fixtures import constants
+        from seleniumbase.fixtures import shared_utils
+
+        pip_find_lock = fasteners.InterProcessLock(
+            constants.PipInstall.FINDLOCK
+        )
+        with pip_find_lock:  # Prevent multi-processes mode issues
+            if not find_spec("boto"):
+                shared_utils.pip_install("boto", version="2.49.0")
+        from boto.s3.connection import S3Connection
 
         self.conn = S3Connection(selenium_access_key, selenium_secret_key)
         self.bucket = self.conn.get_bucket(log_bucket)
         self.bucket_url = bucket_url
 
     def get_key(self, file_name):
         """Create a new Key instance with the given name."""
-        try:
-            from boto.s3.key import Key
-        except Exception:
-            shared_utils.pip_install("boto", version="2.49.0")
-            from boto.s3.key import Key
+        from boto.s3.key import Key
 
         return Key(bucket=self.bucket, name=file_name)