more robusting pickling of saved browser instances

wabot/api.py

@@ -52,7 +52,8 @@ class BrowserProxy:
             'firefox1',
             'firefox2',
             'chromium2',
-            'remote_chromium2',
+            'remote_chromium',
+            'remote_chrome',
             'phantomjs',
             'remote_firefox'
         ), 'webdriver_type must be firefox1, firefox2, chromium2, remote_chromium2, or phantomjs'
@@ -215,7 +216,10 @@ class BrowserProxy:
         if browser == 'chromium2': # Selenium 2 - Chrome
             driver = self._create_driver_chromium()
 
-        elif browser == 'remote_chromium':
+        elif (
+            browser == 'remote_chromium'
+            or browser == 'remote_chrome'
+        ):
             driver = browser_factory._create_driver_remote_chromium(session_name)
 
         elif browser == 'remote_firefox': # Selenium 1 - Firefox

wabot/create_browser.py

@@ -1,5 +1,6 @@
 
 import appdirs
+import datetime
 import logging
 import os
 import selenium
@@ -13,7 +14,17 @@ REFUSE_AFTER_EXCEPTION = True
 EXECUTOR_PORT = 4444
 REMOTE_EXECUTOR = 'http://127.0.0.1:%s/wd/hub'
 
-# Factory Creator 
+
+class PickledBrowserReference:
+    """
+    Structure for saving a webdriver instance. Also includes the timestamp
+    to help invalidate old references.
+    """
+    def __init__(self):
+        self.browser_ref = None
+        self.timestamp = None
+
+# Factory Creator - instantiate Selenium1 and Selenium2 webdriver instances.
 class CreateBrowser:
     """
     Creates and instantiates selenium webbrowser instances.
@@ -29,9 +40,14 @@ class CreateBrowser:
         self._pickle_filename = None
         if pickle_filename is not None:
             self.pickle_filename = pickle_filename
+        self._has_cleaned = False
 
     @property
     def pickle_filename(self):
+        """
+        Reference to the user-data-dir pickle target filename, for
+        saving browser references.
+        """
         if self._pickle_filename is None:
             self._pickle_filename = os.path.join(
                 appdirs.user_data_dir('wabot'),
@@ -45,31 +61,72 @@ class CreateBrowser:
     def pickle_filename(self, value):
         self._pickle_filename = value
 
-    # i think should be chromium1, ie selenium1 or seleniumRC
+    def _clean_old_pickles(self):
-    def _create_driver_remote_chromium(self, session_name):
+        if self._has_cleaned is True:
+            LOGGER.debug('not cleaning pickles twice')
+            return
+
         p = self.pickle_filename
-        final_name = '{}-{}'.format('remote-chromium', session_name)
+        try:
-        driver = None
+            fp = open(p, 'rb')
+            pickles = pickle.load(fp)
+            if not pickles:
+                raise Exception
+
+
+            LOGGER.info('cleaning any old pickle refs')
+
+
+            for pickle_ref in pickles:
+                driver_ref = pickles[pickle_ref]
+                ts = driver_ref.timestamp
+
+                if datetime.datetime.now() - driver_ref.timestamp >= datetime.timedelta(days=3):
+                    LOGGER.info('deleting old pickled driver ref: %s at %s', pickle_ref, ts)
+                    pickles.pop(final_name)
+
+            self._has_cleaned = True
+
+        except (FileNotFoundError, IOError) as ex:
+            LOGGER.error('unable to load pickles: no pickled drivers found')
+        except Exception as ex:
+            print(ex)
+            pass
+
+    def _load_from_pickle(self, final_name):
+        p = self.pickle_filename
+
 
         # Definitely no browser instance already, we must instantiate
         if not os.path.exists(p):
             LOGGER.debug('no pickled file for saved browser instances (nothing saved yet)')
+            return
 
         # There MAY be an open browser or an invalidated reference to a once-open browser
         if os.path.exists(p):
             LOGGER.debug('found pickled file for saved browser instances: %s', p)
             # First, see if existing session_name browser instance exists
             fp = None
-            drivers = {}
+
+            self._clean_old_pickles()
             try:
                 fp = open(p, 'rb')
-                drivers = pickle.load(fp)
+                # drivers = pickle.load(fp)
-                if not drivers:
+                pickles = pickle.load(fp)
+                if not pickles:
                     raise Exception
-                LOGGER.debug('found saved browser instances: %s', list(drivers.keys()))
+
-                driver = drivers.get(final_name)
+                for idx, saved_instance in enumerate(pickles):
-                if not driver:
+                    LOGGER.debug('found saved browser instances (%s): %s', idx, saved_instance)
+
+                driver_ref = pickles.get(final_name)
+                if not driver_ref:
                     raise Exception
+
+                # if datetime.datetime.now() - driver_ref.timestamp >= datetime.timedelta(days=3):
+                #     pickles.pop(final_name)
+
+                driver = driver_ref.browser_ref
                 LOGGER.debug('connected to pickled webdriver instance: %s', final_name)
                 url = driver.current_url # throw error if driver isn't reliable anymore
                 LOGGER.info('webdriver instance is ready')
@@ -78,8 +135,71 @@ class CreateBrowser:
             except (FileNotFoundError, IOError) as ex:
                 LOGGER.error('unable to connect to existing webdriver: no pickled drivers found')
             except Exception as ex:
+                print(ex)
                 self.driver = None
 
+    def _save_to_pickle(self, final_name, driver_ref):
+        p = self.pickle_filename
+        pickles = {}
+
+        # Definitely no browser instance already, we must instantiate
+        if not os.path.exists(p):
+            LOGGER.debug('no pickled file for saved browser instances (nothing saved yet)')
+
+            # Create usr-app-dir if doesn't exist
+            dirname = os.path.dirname(p)
+            if not os.path.isdir(dirname):
+                os.mkdir(dirname)
+
+        # There MAY be an open browser or an invalidated reference to a once-open browser
+        if os.path.exists(p):
+            LOGGER.debug('found existing pickle file for saved browser instances: %s', p)
+            # First, see if existing session_name browser instance exists
+
+
+            self._clean_old_pickles()
+
+            fp = open(p, 'rb')
+            pickles = pickle.load(fp)
+            if not pickles:
+                raise Exception
+            LOGGER.debug('found saved browser instances while saving: %s', list(pickles.keys()))
+            fp.close()
+
+
+        # Save to pickle - creating file if necessary
+        pkl = PickledBrowserReference()
+        pkl.browser_ref = driver_ref
+        pkl.timestamp = datetime.datetime.now()
+        pickles[final_name] = pkl
+
+        
+
+        import pprint;
+        print('pickles:')
+        pprint.pprint(pickles)
+
+        # pickle file must be created
+        fp = open(p, 'wb')
+        # drivers[final_name] = driver
+        LOGGER.info('saving browser instance to pickle: %s', final_name)
+        pickle.dump(pickles, fp)
+        fp.close()
+
+    def _create_driver_remote_chromium(self, session_name):
+        """
+        Creates and returns Selenium1 chromium instance.
+        """
+        p = self.pickle_filename
+        # e.g.: remote-chromium-webdriver
+        final_name = '{}-{}'.format('remote-chromium', session_name)
+        driver = None
+
+        driver = self._load_from_pickle(final_name)
+        if driver is not None:
+            LOGGER.info('webdriver instance is ready (from pickle)')
+            return driver
+
         # At this point, need to instantiate a new browser instance
         sel_host = REMOTE_EXECUTOR % (EXECUTOR_PORT)
         LOGGER.info('instantianting new browser instance (remote_chromium)')
@@ -94,50 +214,25 @@ class CreateBrowser:
             desired_capabilities = opt.to_capabilities()
         )
 
-        # Save to pickle
+        self._save_to_pickle(final_name, driver)
-        fp = open(p, 'wb')
+
-        drivers[final_name] = driver
+        # self.driver = driver
-        LOGGER.info('saving browser instance to pickle: %s', final_name)
-        pickle.dump(drivers, fp)
-        self.driver = driver
         return driver
 
 
     # Pickle impl. is duped here
 
     def _create_driver_remote_firefox(self, session_name):
-        p = self.pickle_filename
+        """
+        Creates and returns Selenium1 firefox instance.
+        """
         final_name = '{}-{}'.format('remote-firefox', session_name)
         driver = None
 
-        # Definitely no browser instance already, we must instantiate
+        driver = self._load_from_pickle(final_name)
-        if not os.path.exists(p):
+        if driver is not None:
-            LOGGER.debug('no pickled file for saved browser instances (nothing saved yet)')
+            LOGGER.info('webdriver instance is ready (from pickle)')
-
+            return driver
-        # There MAY be an open browser or an invalidated reference to a once-open browser
-        if os.path.exists(p):
-            LOGGER.debug('found pickled file for saved browser instances: %s', p)
-            # First, see if existing session_name browser instance exists
-            fp = None
-            drivers = {}
-            try:
-                fp = open(p, 'rb')
-                drivers = pickle.load(fp)
-                if not drivers:
-                    raise Exception
-                LOGGER.debug('found saved browser instances: %s', list(drivers.keys()))
-                driver = drivers.get(final_name)
-                if not driver:
-                    raise Exception
-                LOGGER.debug('connected to pickled webdriver instance: %s', final_name)
-                url = driver.current_url # throw error if driver isn't reliable anymore
-                LOGGER.info('webdriver instance is ready')
-                # self.driver = driver
-                return driver
-            except (FileNotFoundError, IOError) as ex:
-                LOGGER.error('unable to connect to existing webdriver: no pickled drivers found')
-            except Exception as ex:
-                self.driver = None
 
         # At this point, need to instantiate a new browser instance
         sel_host = REMOTE_EXECUTOR % (EXECUTOR_PORT)
@@ -153,12 +248,7 @@ class CreateBrowser:
             browser_profile = profile
         )
 
-        # Save to pickle
+        self._save_to_pickle(final_name, driver)
-        fp = open(p, 'wb')
-        drivers[final_name] = driver
-        LOGGER.info('saving browser instance to pickle: %s', final_name)
-        pickle.dump(drivers, fp)
-
         return driver