image_download/classes.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import logging
 
 from selenium import webdriver
+from selenium.common.exceptions import TimeoutException
 from selenium.webdriver.chrome.service import Service
 from selenium.webdriver.common.action_chains import ActionChains
 from selenium.webdriver.common.by import By
@@ -49,7 +50,7 @@ class Instagram:
                 ),
             )
             logging.debug('located button')
-        except TimeoutError:
+        except TimeoutException:
             logging.error('Timeout before cookie popup appeared.')
         else:
             elem.click()
@@ -58,18 +59,18 @@ class Instagram:
         img_elem_list = self.webdriver.find_elements(
             by=By.XPATH,
             value=(
-                "//article[@role='presentation']"
-                "//div[@role='presentation']"
+                "//*[@role='presentation']"
                 "//img[contains(@src,'https://scontent.cdninstagram.com/v/')]"
             ),
         )
+        assert len(img_elem_list) > 0
         # extract src url
         img_src_list = [
             img_elem.get_attribute(
-                'src',
+                name='src',
             ) for img_elem in img_elem_list
         ]
-        # remove src urls with size mathing profile images
+        # remove src urls with size matching profile images
         img_src_list = [
             img_src for img_src in img_src_list if 's150x150' not in img_src
         ]