class AWS::SimpleWorkflow::ActivityTaskCollection

def count task_list
  options = {}
  options[:domain] = domain.name
  options[:task_list] = { :name => task_list }
  response = client.count_pending_activity_tasks(options)
  Count.new(response.data['count'], response.data['truncated'])
end

def initialize domain, options = {}
  @domain = domain
  super
end

def poll task_list, options = {}, &block
  loop do
    begin
      poll_for_single_task(task_list, options) do |activity_task|
        yield(activity_task)
      end
    rescue Timeout::Error
      retry
    end
  end
  nil
end

def poll_for_single_task task_list, options = {}, &block
  client_opts = {}
  client_opts[:domain] = domain.name
  client_opts[:task_list] = { :name => task_list }
  client_opts[:identity] = identity_opt(options)
  response = client.poll_for_activity_task(client_opts)
  if response.data['taskToken']
    activity_task = ActivityTask.new(domain, response.data)
    if block_given?
      begin
        yield(activity_task)
        activity_task.complete! unless activity_task.responded?
      rescue ActivityTask::CancelRequestedError
        activity_task.cancel! unless activity_task.responded?
      rescue StandardError => e
        unless activity_task.responded?
          reason = "UNTRAPPED ERROR: #{e.message}"
          details = e.backtrace.join("\n")
          activity_task.fail!(:reason => reason, :details => details)
        end
        raise e
      end
      nil
    else
      activity_task
    end
  else
    nil
  end
end