returned the iterator and corrected the test

Author: mishushakov

browserbase/__init__.py

@@ -37,18 +37,13 @@ def load_urls(self, urls: List[str]):
                 "wss://api.browserbase.com?apiKey=" + self.api_key
             )
 
-            html_contents = []
             for url in urls:
                 page = browser.new_page()
                 page.goto(url)
-                html = page.content()
-                html_contents.append(html)
-                page.close()
+                yield page.content()
 
             browser.close()
 
-            return html_contents
-
     def screenshot(self, url: str, full_page: bool = False):
         """Load a page in a headless browser and return a screenshot as base64 string"""
         if not url:

tests/lib.py

@@ -12,13 +12,12 @@ def test_load(self):
         self.assertIn("Example Domain", result)
 
     def test_load_urls(self):
-        result = self.browserbase.load_urls(["https://example.com"])
-        self.assertIn("Example Domain", result[0])
+        result = next(self.browserbase.load_urls(["https://example.com"]))
+        self.assertIn("Example Domain", result)
 
     def test_screenshot(self):
         result = self.browserbase.screenshot("https://example.com")
-        print(result)
-        self.assertIn("Example Domain", result)
+        self.assertIn("oq0IBiQAAAABJRU5ErkJggg==", result)
 
 
 if __name__ == "__main__":