add spec for BrowserContext, and fix some bugs

YusukeIwaki · Jul 24, 2020 · ac314aa · ac314aa
1 parent 277183b
commit ac314aa
Show file tree

Hide file tree

Showing 8 changed files with 247 additions and 18 deletions.
diff --git a/lib/puppeteer.rb b/lib/puppeteer.rb
@@ -171,7 +171,7 @@ def connect(
       begin
         yield(browser)
       ensure
-        browser.close
+        browser.disconnect
       end
     else
       browser

diff --git a/lib/puppeteer/browser.rb b/lib/puppeteer/browser.rb
@@ -43,7 +43,7 @@ def initialize(connection:, context_ids:, ignore_https_errors:, default_viewport
     @default_context = Puppeteer::BrowserContext.new(@connection, self, nil)
     @contexts = {}
     context_ids.each do |context_id|
-      @contexts[context_id] = Puppeteer::BrowserContext.new(@connection, self. context_id)
+      @contexts[context_id] = Puppeteer::BrowserContext.new(@connection, self, context_id)
     end
     @targets = {}
     @connection.on_event 'Events.Connection.Disconnected' do
@@ -70,6 +70,15 @@ def on(event_name, &block)
     add_event_listener(EVENT_MAPPINGS[event_name.to_sym], &block)
   end
 
+  # @param event_name [Symbol]
+  def once(event_name, &block)
+    unless EVENT_MAPPINGS.has_key?(event_name.to_sym)
+      raise ArgumentError.new("Unknown event name: #{event_name}. Known events are #{EVENT_MAPPINGS.keys.join(", ")}")
+    end
+
+    observe_first(EVENT_MAPPINGS[event_name.to_sym], &block)
+  end
+
   # @return [Puppeteer::BrowserRunner::BrowserProcess]
   def process
     @process
@@ -94,7 +103,7 @@ def default_browser_context
   # @param context_id [String]
   def dispose_context(context_id)
     @connection.send_message('Target.disposeBrowserContext', browserContextId: context_id)
-    @contexts.remove(context_id)
+    @contexts.delete(context_id)
   end
 
   class TargetAlreadyExistError < StandardError
@@ -204,12 +213,11 @@ def target
     @targets[target_id]
   end
 
-  # @param {function(!Target):boolean} predicate
-  # @param {{timeout?: number}=} options
-  # @return {!Promise<!Target>}
+  # @param predicate [Proc(Puppeteer::Target -> Boolean)]
+  # @return [Puppeteer::Target]
   def wait_for_target(predicate:, timeout: nil)
     timeout_in_sec = (timeout || 30000).to_i / 1000.0
-    existing_target = targets.first { |target| predicate.call(target) }
+    existing_target = targets.find { |target| predicate.call(target) }
     return existing_target if existing_target
 
     event_listening_ids = []
@@ -233,11 +241,18 @@ def wait_for_target(predicate:, timeout: nil)
       else
         target_promise.value!
       end
+    rescue Timeout::Error
+      raise Puppeteer::TimeoutError.new("waiting for target failed: timeout #{timeout}ms exceeded")
     ensure
       remove_event_listener(*event_listening_ids)
     end
   end
 
+  # @!method async_wait_for_target(predicate:, timeout: nil)
+  #
+  # @param predicate [Proc(Puppeteer::Target -> Boolean)]
+  define_async_method :async_wait_for_target
+
   # @return {!Promise<!Array<!Puppeteer.Page>>}
   def pages
     browser_contexts.flat_map(&:pages)

diff --git a/lib/puppeteer/browser_context.rb b/lib/puppeteer/browser_context.rb
@@ -1,5 +1,6 @@
 class Puppeteer::BrowserContext
   include Puppeteer::EventCallbackable
+  using Puppeteer::DefineAsyncMethod
 
   # @param {!Puppeteer.Connection} connection
   # @param {!Browser} browser
@@ -10,28 +11,57 @@ def initialize(connection, browser, context_id)
     @id = context_id
   end
 
+  EVENT_MAPPINGS = {
+    disconnected: 'Events.BrowserContext.Disconnected',
+    targetcreated: 'Events.BrowserContext.TargetCreated',
+    targetchanged: 'Events.BrowserContext.TargetChanged',
+    targetdestroyed: 'Events.BrowserContext.TargetDestroyed',
+  }
+
+  # @param event_name [Symbol] either of :disconnected, :targetcreated, :targetchanged, :targetdestroyed
+  def on(event_name, &block)
+    unless EVENT_MAPPINGS.has_key?(event_name.to_sym)
+      raise ArgumentError.new("Unknown event name: #{event_name}. Known events are #{EVENT_MAPPINGS.keys.join(", ")}")
+    end
+
+    add_event_listener(EVENT_MAPPINGS[event_name.to_sym], &block)
+  end
+
+  # @param event_name [Symbol]
+  def once(event_name, &block)
+    unless EVENT_MAPPINGS.has_key?(event_name.to_sym)
+      raise ArgumentError.new("Unknown event name: #{event_name}. Known events are #{EVENT_MAPPINGS.keys.join(", ")}")
+    end
+
+    observe_first(EVENT_MAPPINGS[event_name.to_sym], &block)
+  end
+
   # @return {!Array<!Target>} target
   def targets
     @browser.targets.select { |target| target.browser_context == self }
   end
 
-  # @param {function(!Target):boolean} predicate
-  # @param {{timeout?: number}=} options
-  # @return {!Promise<!Target>}
+  # @param predicate [Proc(Puppeteer::Target -> Boolean)]
+  # @return [Puppeteer::Target]
   def wait_for_target(predicate:, timeout: nil)
     @browser.wait_for_target(
       predicate: ->(target) { target.browser_context == self && predicate.call(target) },
       timeout: timeout,
     )
   end
 
+  # @!method async_wait_for_target(predicate:, timeout: nil)
+  #
+  # @param predicate [Proc(Puppeteer::Target -> Boolean)]
+  define_async_method :async_wait_for_target
+
   # @return {!Promise<!Array<!Puppeteer.Page>>}
   def pages
     targets.select { |target| target.type == 'page' }.map(&:page).reject { |page| !page }
   end
 
   def incognito?
-    !@id
+    !!@id
   end
 
   # /**
@@ -82,7 +112,7 @@ def browser
   end
 
   def close
-    if !@id
+    unless @id
       raise 'Non-incognito profiles cannot be closed!'
     end
     @browser.dispose_context(@id)

diff --git a/lib/puppeteer/connection.rb b/lib/puppeteer/connection.rb
@@ -49,7 +49,7 @@ def initialize(url, transport, delay = 0)
       async_handle_message(message)
     end
     @transport.on_close do |reason, code|
-      handle_close(reason, code)
+      handle_close
     end
 
     @sessions = {}

diff --git a/lib/puppeteer/page.rb b/lib/puppeteer/page.rb
@@ -1080,8 +1080,12 @@ def select(selector, *values)
   define_async_method :async_select
 
   # @param selector [String]
-  def tap(selector)
-    main_frame.tap(selector)
+  def tap(selector: nil, &block)
+    if selector.nil? && block
+      super(&block)
+    else
+      main_frame.tap(selector)
+    end
   end
 
   define_async_method :async_tap

diff --git a/lib/puppeteer/web_socket.rb b/lib/puppeteer/web_socket.rb
@@ -16,10 +16,14 @@ def initialize(url)
 
     def write(data)
       @socket.write(data)
+    rescue Errno::EPIPE
+      raise EOFError.new('already closed')
     end
 
     def readpartial(maxlen = 1024)
       @socket.readpartial(maxlen)
+    rescue Errno::ECONNRESET
+      raise EOFError.new('closed by remote')
     end
   end
 
@@ -40,6 +44,9 @@ def initialize(url:, max_payload_size:)
     rescue EOFError
       # Google Chrome was gone.
       # We have nothing todo. Just finish polling.
+      if @ready_state < STATE_CLOSING
+        handle_on_close(reason: 'Going Away', code: 1001)
+      end
     end
   end
 

diff --git a/spec/integration/browser_context_spec.rb b/spec/integration/browser_context_spec.rb
@@ -0,0 +1,173 @@
+require 'spec_helper'
+
+RSpec.describe Puppeteer::BrowserContext, puppeteer: :browser do
+  describe 'default context' do
+    it 'should have default context' do
+      expect(browser.browser_contexts.length).to eq(1)
+
+      default_context = browser.browser_contexts.first
+      expect(default_context).not_to be_incognito
+    end
+
+    it 'cannot be closed' do
+      default_context = browser.browser_contexts.first
+      expect { default_context.close }.to raise_error(/cannot be closed/)
+    end
+  end
+
+  describe 'create incognito context' do
+    it 'should create new incognito context' do
+      context = browser.create_incognito_browser_context
+      expect(context).to be_incognito
+      expect(browser.browser_contexts.length).to eq(2)
+      expect(browser.browser_contexts).to include(context)
+
+      context.close
+      expect(browser.browser_contexts.length).to eq(1)
+    end
+
+    it 'should close all belonging targets once closing context' do
+      expect(browser.pages.length).to eq(1)
+
+      context = browser.create_incognito_browser_context
+      page = context.new_page
+      expect(browser.pages.length).to eq(2)
+      expect(context.pages.length).to eq(1)
+
+      context.close
+      expect(browser.pages.length).to eq(1)
+    end
+
+    it 'window.open should use parent tab context' do
+      context = browser.create_incognito_browser_context
+      page = context.new_page
+      page.goto('about:blank')
+
+      popup_target = await_all(
+        resolvable_future { |f| browser.once('targetcreated') { |target| f.fulfill(target) } },
+        page.async_evaluate('url => { window.open(url); return null }', 'about:blank')
+      ).first
+      expect(popup_target.browser_context).to eq(context)
+      context.close
+    end
+  end
+
+  describe 'target events' do
+    sinatra do
+      get '/test' do
+        'test'
+      end
+    end
+
+    it 'should fire target events' do
+      context = browser.create_incognito_browser_context
+      events = []
+      context.on('targetcreated') do |target|
+        events << "CREATED: #{target.url}"
+      end
+      context.on('targetchanged') do |target|
+        events << "CHANGED: #{target.url}"
+      end
+      context.on('targetdestroyed') do |target|
+        events << "DESTROYED: #{target.url}"
+      end
+
+      page = context.new_page
+      page.goto('http://127.0.0.1:4567/test')
+      page.close
+
+      expect(events).to eq([
+        "CREATED: about:blank",
+        "CHANGED: http://127.0.0.1:4567/test",
+        "DESTROYED: http://127.0.0.1:4567/test",
+      ])
+      context.close
+    end
+  end
+
+  describe 'wait for target' do
+    sinatra do
+      get '/test' do
+        'test'
+      end
+    end
+
+    it 'should wait for a target' do
+      context = browser.create_incognito_browser_context
+      resolved = false
+      target_promise = context.async_wait_for_target(predicate: -> (target) { target.url == 'http://127.0.0.1:4567/test' })
+      target_promise.then { resolved = true }
+
+      page = context.new_page
+      expect(resolved).to eq(false)
+      page.goto('http://127.0.0.1:4567/test')
+      target = await target_promise
+      expect(target.page).to eq(page)
+      context.close
+    end
+
+    it 'should timeout waiting for a non-existent target' do
+      context = browser.create_incognito_browser_context
+      resolved = false
+      target_promise = context.async_wait_for_target(timeout: 500, predicate: -> (target) { target.url == '?????' })
+      target_promise.then { resolved = true }
+
+      page = context.new_page
+      expect(resolved).to eq(false)
+      page.goto('http://127.0.0.1:4567/test')
+      expect(resolved).to eq(false)
+      expect { await target_promise }.to raise_error(Puppeteer::TimeoutError)
+      context.close
+    end
+  end
+
+  describe 'isolation' do
+    sinatra do
+      get '/isolation' do
+        'test isolation'
+      end
+    end
+
+    it 'should isolate localStorage and cookies' do
+      # Create two incognito contexts.
+      contexts = 2.times.map { browser.create_incognito_browser_context }
+
+      contexts.each do |context|
+        expect(context.targets.length).to eq(0)
+      end
+
+      pages = contexts.map.with_index do |context, index|
+        context.new_page.tap do |page|
+          page.goto('http://127.0.0.1:4567/isolation')
+          page.evaluate <<~JAVASCRIPT
+          () => {
+            localStorage.setItem('name', 'page#{index}');
+            document.cookie = 'name=page#{index}';
+          }
+          JAVASCRIPT
+        end
+      end
+
+      contexts.each_with_index do |context, index|
+        expect(context.targets.length).to eq(1)
+
+        # Make sure pages don't share localstorage or cookies.
+        expect(pages[index].evaluate("() => localStorage.getItem('name')")).to eq("page#{index}")
+        expect(pages[index].evaluate("() => document.cookie")).to eq("name=page#{index}")
+      end
+
+      contexts.each(&:close)
+      expect(browser.browser_contexts.length).to eq(1)
+    end
+
+    it 'should work across sessions' do
+      expect(browser.browser_contexts.length).to eq(1)
+      context = browser.create_incognito_browser_context
+      expect(browser.browser_contexts.length).to eq(2)
+      Puppeteer.connect(browser_ws_endpoint: browser.ws_endpoint) do |remote_browser|
+        expect(remote_browser.browser_contexts.length).to eq(2)
+      end
+      context.close
+    end
+  end
+end
diff --git a/spec/integration/browser_spec.rb b/spec/integration/browser_spec.rb
@@ -23,10 +23,10 @@
     it 'should return the browser connected state' do
       ws_endpoint = browser.ws_endpoint
       new_browser = Puppeteer.connect(browser_ws_endpoint: ws_endpoint)
-      expect(new_browser.connected?).to eq(true)
+      expect(new_browser).to be_connected
       new_browser.disconnect
-      expect(new_browser.connected?).to eq(false)
-      expect(browser.connected?).to eq(true)
+      expect(new_browser).not_to be_connected
+      expect(browser).to be_connected
     end
   end
 end
-Original file line number
+Diff line change
@@ Expand Up / @@ -171,7 +171,7 @@ def connect( @@
           begin
             yield(browser)
           ensure
-            browser.close
+            browser.disconnect
           end
         else
           browser
@@ Expand Down @@