fluent-plugin-bigquery 0.2.16 → 0.3.0

data/test/helper.rb

@@ -29,6 +29,7 @@ require 'fluent/plugin/buf_file'
 require 'fluent/plugin/out_bigquery'
 
 require 'rr'
+require 'test/unit/rr'
 
 class Test::Unit::TestCase
 end

data/test/plugin/test_out_bigquery.rb

@@ -34,20 +34,10 @@ class BigQueryOutputTest < Test::Unit::TestCase
     Fluent::Test::TimeSlicedOutputTestDriver.new(Fluent::BigQueryOutput).configure(conf)
   end
 
-  def stub_client(driver)
-    stub(client = Object.new) do |expect|
-      yield expect if defined?(yield)
-    end
-    stub(driver.instance).client { client }
-    client
-  end
-
-  def mock_client(driver)
-    mock(client = Object.new) do |expect|
-      yield expect
-    end
-    stub(driver.instance).client { client }
-    client
+  def stub_writer(driver)
+    writer = driver.instance.writer
+    stub(writer).get_auth { nil }
+    writer
   end
 
   def test_configure_table
@@ -76,21 +66,25 @@ class BigQueryOutputTest < Test::Unit::TestCase
       issuer: 'foo@bar.example',
       signing_key: key) { authorization }
 
-    mock.proxy(Google::Apis::BigqueryV2::BigqueryService).new.with_any_args {
-      mock!.__send__(:authorization=, authorization) {}
-    }
+    mock.proxy(Google::Apis::BigqueryV2::BigqueryService).new.with_any_args do |cl|
+      mock(cl).__send__(:authorization=, authorization) {}
+      cl
+    end
 
-    driver = create_driver(CONFIG)
-    driver.instance.client()
+    driver = create_driver
+    mock.proxy(Fluent::BigQuery::Writer).new(duck_type(:info, :error, :warn), driver.instance.auth_method, is_a(Hash))
+    driver.instance.writer
+    assert driver.instance.writer.client.is_a?(Google::Apis::BigqueryV2::BigqueryService)
   end
 
   def test_configure_auth_compute_engine
     authorization = Object.new
     mock(Google::Auth::GCECredentials).new { authorization }
 
-    mock.proxy(Google::Apis::BigqueryV2::BigqueryService).new.with_any_args {
-      mock!.__send__(:authorization=, authorization) {}
-    }
+    mock.proxy(Google::Apis::BigqueryV2::BigqueryService).new.with_any_args do |cl|
+      mock(cl).__send__(:authorization=, authorization) {}
+      cl
+    end
 
     driver = create_driver(%[
       table foo
@@ -99,7 +93,9 @@ class BigQueryOutputTest < Test::Unit::TestCase
       dataset yourdataset_id
       field_integer time,status,bytes
     ])
-    driver.instance.client()
+    mock.proxy(Fluent::BigQuery::Writer).new(duck_type(:info, :error, :warn), driver.instance.auth_method, is_a(Hash))
+    driver.instance.writer
+    assert driver.instance.writer.client.is_a?(Google::Apis::BigqueryV2::BigqueryService)
   end
 
   def test_configure_auth_json_key_as_file
@@ -107,9 +103,10 @@ class BigQueryOutputTest < Test::Unit::TestCase
     authorization = Object.new
     mock(Google::Auth::ServiceAccountCredentials).make_creds(json_key_io: File.open(json_key_path), scope: API_SCOPE) { authorization }
 
-    mock.proxy(Google::Apis::BigqueryV2::BigqueryService).new.with_any_args {
-      mock!.__send__(:authorization=, authorization) {}
-    }
+    mock.proxy(Google::Apis::BigqueryV2::BigqueryService).new.with_any_args do |cl|
+      mock(cl).__send__(:authorization=, authorization) {}
+      cl
+    end
 
     driver = create_driver(%[
       table foo
@@ -119,7 +116,32 @@ class BigQueryOutputTest < Test::Unit::TestCase
       dataset yourdataset_id
       field_integer time,status,bytes
     ])
-    driver.instance.client()
+    mock.proxy(Fluent::BigQuery::Writer).new(duck_type(:info, :error, :warn), driver.instance.auth_method, is_a(Hash))
+    driver.instance.writer
+    assert driver.instance.writer.client.is_a?(Google::Apis::BigqueryV2::BigqueryService)
+  end
+
+  def test_configure_auth_json_key_as_file_raise_permission_error
+    json_key_path = 'test/plugin/testdata/json_key.json'
+    json_key_path_dir = File.dirname(json_key_path)
+
+    begin
+      File.chmod(0000, json_key_path_dir)
+
+      driver = create_driver(%[
+        table foo
+        auth_method json_key
+        json_key #{json_key_path}
+        project yourproject_id
+        dataset yourdataset_id
+        field_integer time,status,bytes
+      ])
+      assert_raises(Errno::EACCES) do
+        driver.instance.writer.client
+      end
+    ensure
+      File.chmod(0755, json_key_path_dir)
+    end
   end
 
   def test_configure_auth_json_key_as_string
@@ -129,9 +151,10 @@ class BigQueryOutputTest < Test::Unit::TestCase
     authorization = Object.new
     mock(Google::Auth::ServiceAccountCredentials).make_creds(json_key_io: json_key_io, scope: API_SCOPE) { authorization }
 
-    mock.proxy(Google::Apis::BigqueryV2::BigqueryService).new.with_any_args {
-      mock!.__send__(:authorization=, authorization) {}
-    }
+    mock.proxy(Google::Apis::BigqueryV2::BigqueryService).new.with_any_args do |cl|
+      mock(cl).__send__(:authorization=, authorization) {}
+      cl
+    end
 
     driver = create_driver(%[
       table foo
@@ -141,16 +164,19 @@ class BigQueryOutputTest < Test::Unit::TestCase
       dataset yourdataset_id
       field_integer time,status,bytes
     ])
-    driver.instance.client()
+    mock.proxy(Fluent::BigQuery::Writer).new(duck_type(:info, :error, :warn), driver.instance.auth_method, is_a(Hash))
+    driver.instance.writer
+    assert driver.instance.writer.client.is_a?(Google::Apis::BigqueryV2::BigqueryService)
   end
 
   def test_configure_auth_application_default
     authorization = Object.new
     mock(Google::Auth).get_application_default([API_SCOPE]) { authorization }
 
-    mock.proxy(Google::Apis::BigqueryV2::BigqueryService).new.with_any_args {
-      mock!.__send__(:authorization=, authorization) {}
-    }
+    mock.proxy(Google::Apis::BigqueryV2::BigqueryService).new.with_any_args do |cl|
+      mock(cl).__send__(:authorization=, authorization) {}
+      cl
+    end
 
     driver = create_driver(%[
       table foo
@@ -159,7 +185,10 @@ class BigQueryOutputTest < Test::Unit::TestCase
       dataset yourdataset_id
       field_integer time,status,bytes
     ])
-    driver.instance.client()
+
+    mock.proxy(Fluent::BigQuery::Writer).new(duck_type(:info, :error, :warn), driver.instance.auth_method, is_a(Hash))
+    driver.instance.writer
+    assert driver.instance.writer.client.is_a?(Google::Apis::BigqueryV2::BigqueryService)
   end
 
   def test_configure_fieldname_stripped
@@ -275,9 +304,15 @@ class BigQueryOutputTest < Test::Unit::TestCase
         "requesttime" => (now - 1).to_f.to_s.to_f,
         "bot_access" => true,
         "loginsession" => false,
+        "something-else" => "would be ignored",
+        "yet-another" => {
+          "foo" => "bar",
+          "baz" => 1,
+        },
         "remote" => {
           "host" => "remote.example",
           "ip" =>  "192.0.2.1",
+          "port" => 12345,
           "user" => "tagomoris",
         }
       }
@@ -327,7 +362,6 @@ class BigQueryOutputTest < Test::Unit::TestCase
         time_field  time
         #{type}     time
       CONFIG
-      stub_client(driver)
 
       driver.instance.start
       buf = driver.instance.format_stream("my.tag", [input])
@@ -371,7 +405,7 @@ class BigQueryOutputTest < Test::Unit::TestCase
       field_integer metadata.time
       field_string  metadata.node,log
     CONFIG
-    stub_client(driver)
+
     driver.instance.start
     buf = driver.instance.format_stream("my.tag", [input])
     driver.instance.shutdown
@@ -429,12 +463,18 @@ class BigQueryOutputTest < Test::Unit::TestCase
         "remote" => {
           "host" => "remote.example",
           "ip" =>  "192.0.2.1",
+          "port" => 12345,
           "user" => "tagomoris",
         },
         "response" => {
           "status" => 1,
           "bytes" => 3,
         },
+        "something-else" => "would be ignored",
+        "yet-another" => {
+          "foo" => "bar",
+          "baz" => 1,
+        },
       }
     }
 
@@ -531,16 +571,10 @@ class BigQueryOutputTest < Test::Unit::TestCase
       fetch_schema true
       field_integer time
     CONFIG
-    mock_client(driver) do |expect|
-      expect.get_table('yourproject_id', 'yourdataset_id', 'foo') {
-        s = stub!
-        schema_stub = stub!
-        fields_stub = stub!
-        s.schema { schema_stub }
-        schema_stub.fields { fields_stub }
-        fields_stub.as_json { sudo_schema_response.deep_stringify_keys["schema"]["fields"] }
-        s
-      }
+
+    writer = stub_writer(driver)
+    mock(writer).fetch_schema('yourproject_id', 'yourdataset_id', 'foo') do
+      sudo_schema_response.deep_stringify_keys["schema"]["fields"]
     end
     driver.instance.start
     buf = driver.instance.format_stream("my.tag", [input])
@@ -603,16 +637,10 @@ class BigQueryOutputTest < Test::Unit::TestCase
       fetch_schema true
       field_integer time
     CONFIG
-    mock_client(driver) do |expect|
-      expect.get_table('yourproject_id', 'yourdataset_id', now.strftime('foo_%Y_%m_%d')) {
-        s = stub!
-        schema_stub = stub!
-        fields_stub = stub!
-        s.schema { schema_stub }
-        schema_stub.fields { fields_stub }
-        fields_stub.as_json { sudo_schema_response.deep_stringify_keys["schema"]["fields"] }
-        s
-      }
+
+    writer = stub_writer(driver)
+    mock(writer).fetch_schema('yourproject_id', 'yourdataset_id', now.strftime('foo_%Y_%m_%d')) do
+      sudo_schema_response.deep_stringify_keys["schema"]["fields"]
     end
     driver.instance.start
     buf = driver.instance.format_stream("my.tag", [input])
@@ -731,6 +759,8 @@ class BigQueryOutputTest < Test::Unit::TestCase
       dataset yourdataset_id
 
       field_string uuid
+
+      buffer_type memory
     CONFIG
     driver.instance.start
     buf = driver.instance.format_stream("my.tag", [input])
@@ -739,38 +769,6 @@ class BigQueryOutputTest < Test::Unit::TestCase
     assert_equal expected, buf
   end
 
-  def test_empty_value_in_required
-    now = Time.now
-    input = [
-      now,
-      {
-        "tty" => "pts/1",
-        "pwd" => "/home/yugui",
-        "user" => nil,
-        "argv" => %w[ tail -f /var/log/fluentd/fluentd.log ]
-      }
-    ]
-
-    driver = create_driver(<<-CONFIG)
-      table foo
-      email foo@bar.example
-      private_key_path /path/to/key
-      project yourproject_id
-      dataset yourdataset_id
-
-      time_format %s
-      time_field  time
-
-      schema_path #{File.join(File.dirname(__FILE__), "testdata", "sudo.schema")}
-      field_integer time
-    CONFIG
-    driver.instance.start
-    assert_raises(RuntimeError.new("Required field user cannot be null")) do
-      driver.instance.format_stream("my.tag", [input])
-    end
-    driver.instance.shutdown
-  end
-
   def test_replace_record_key
     now = Time.now
     input = [
@@ -869,13 +867,119 @@ class BigQueryOutputTest < Test::Unit::TestCase
 
   def test_write
     entry = {json: {a: "b"}}, {json: {b: "c"}}
-    driver = create_driver(CONFIG)
-    mock_client(driver) do |expect|
-      expect.insert_all_table_data('yourproject_id', 'yourdataset_id', 'foo', {
-        rows: entry
-      }, {}) { stub! }
+    driver = create_driver
+
+    writer = stub_writer(driver)
+    mock.proxy(writer).insert_rows('yourproject_id', 'yourdataset_id', 'foo', entry, hash_including(
+      skip_invalid_rows: false,
+      ignore_unknown_values: false
+    ))
+    mock(writer.client).insert_all_table_data('yourproject_id', 'yourdataset_id', 'foo', {
+      rows: entry,
+      skip_invalid_rows: false,
+      ignore_unknown_values: false
+    }, {options: {timeout_sec: nil, open_timeout_sec: 60}}) do
+      s = stub!
+      s.insert_errors { nil }
+      s
+    end
+
+    chunk = Fluent::MemoryBufferChunk.new("my.tag")
+    entry.each do |e|
+      chunk << e.to_msgpack
+    end
+
+    driver.instance.start
+    driver.instance.write(chunk)
+    driver.instance.shutdown
+  end
+
+  def test_write_with_retryable_error
+    entry = {json: {a: "b"}}, {json: {b: "c"}}
+    driver = create_driver(<<-CONFIG)
+      table foo
+      email foo@bar.example
+      private_key_path /path/to/key
+      project yourproject_id
+      dataset yourdataset_id
+
+      time_format %s
+      time_field  time
+
+      field_integer time,status,bytes
+      field_string  vhost,path,method,protocol,agent,referer,remote.host,remote.ip,remote.user
+      field_float   requesttime
+      field_boolean bot_access,loginsession
+      <secondary>
+        type file
+        path error
+        utc
+      </secondary>
+    CONFIG
+
+    writer = stub_writer(driver)
+    mock(writer.client).insert_all_table_data('yourproject_id', 'yourdataset_id', 'foo', {
+      rows: entry,
+      skip_invalid_rows: false,
+      ignore_unknown_values: false
+    }, {options: {timeout_sec: nil, open_timeout_sec: 60}}) do
+      ex = Google::Apis::ServerError.new("error")
+      def ex.reason
+        "backendError"
+      end
+      raise ex
+    end
+
+    chunk = Fluent::MemoryBufferChunk.new("my.tag")
+    entry.each do |e|
+      chunk << e.to_msgpack
+    end
+
+    driver.instance.start
+    assert_raise Fluent::BigQuery::Writer::RetryableError do
+      driver.instance.write(chunk)
+    end
+    driver.instance.shutdown
+  end
+
+  def test_write_with_not_retryable_error
+    entry = {json: {a: "b"}}, {json: {b: "c"}}
+    driver = create_driver(<<-CONFIG)
+      table foo
+      email foo@bar.example
+      private_key_path /path/to/key
+      project yourproject_id
+      dataset yourdataset_id
+
+      time_format %s
+      time_field  time
+
+      field_integer time,status,bytes
+      field_string  vhost,path,method,protocol,agent,referer,remote.host,remote.ip,remote.user
+      field_float   requesttime
+      field_boolean bot_access,loginsession
+      <secondary>
+        type file
+        path error
+        utc
+      </secondary>
+    CONFIG
+
+    writer = stub_writer(driver)
+    mock(writer.client).insert_all_table_data('yourproject_id', 'yourdataset_id', 'foo', {
+      rows: entry,
+      skip_invalid_rows: false,
+      ignore_unknown_values: false
+    }, {options: {timeout_sec: nil, open_timeout_sec: 60}}) do
+      ex = Google::Apis::ServerError.new("error")
+      def ex.reason
+        "invalid"
+      end
+      raise ex
     end
 
+    mock(driver.instance).flush_secondary(is_a(Fluent::Output))
+
     chunk = Fluent::MemoryBufferChunk.new("my.tag")
     entry.each do |e|
       chunk << e.to_msgpack
@@ -902,41 +1006,271 @@ class BigQueryOutputTest < Test::Unit::TestCase
 
       schema_path #{schema_path}
       field_integer time
+
+      buffer_type memory
     CONFIG
     schema_fields = MultiJson.load(File.read(schema_path)).map(&:deep_symbolize_keys).tap do |h|
       h[0][:type] = "INTEGER"
       h[0][:mode] = "NULLABLE"
     end
 
+    writer = stub_writer(driver)
     chunk = Fluent::MemoryBufferChunk.new("my.tag")
     io = StringIO.new("hello")
     mock(driver.instance).create_upload_source(chunk).yields(io)
-    mock_client(driver) do |expect|
-      expect.insert_job('yourproject_id', {
-        configuration: {
-          load: {
-            destination_table: {
-              project_id: 'yourproject_id',
-              dataset_id: 'yourdataset_id',
-              table_id: 'foo',
-            },
-            schema: {
-              fields: schema_fields,
-            },
-            write_disposition: "WRITE_APPEND",
-            source_format: "NEWLINE_DELIMITED_JSON"
-          }
+    mock(writer).wait_load_job("yourproject_id", "yourdataset_id", "dummy_job_id", "foo") { nil }
+    mock(writer.client).insert_job('yourproject_id', {
+      configuration: {
+        load: {
+          destination_table: {
+            project_id: 'yourproject_id',
+            dataset_id: 'yourdataset_id',
+            table_id: 'foo',
+          },
+          schema: {
+            fields: schema_fields,
+          },
+          write_disposition: "WRITE_APPEND",
+          source_format: "NEWLINE_DELIMITED_JSON",
+          ignore_unknown_values: false,
+          max_bad_records: 0,
+        }
+      }
+    }, {upload_source: io, content_type: "application/octet-stream", options: {timeout_sec: nil, open_timeout_sec: 60}}) do
+      s = stub!
+      job_reference_stub = stub!
+      s.job_reference { job_reference_stub }
+      job_reference_stub.job_id { "dummy_job_id" }
+      s
+    end
+
+    entry.each do |e|
+      chunk << MultiJson.dump(e) + "\n"
+    end
+
+    driver.instance.start
+    driver.instance.write(chunk)
+    driver.instance.shutdown
+  end
+
+  def test_write_for_load_with_prevent_duplicate_load
+    schema_path = File.join(File.dirname(__FILE__), "testdata", "sudo.schema")
+    entry = {a: "b"}, {b: "c"}
+    driver = create_driver(<<-CONFIG)
+      method load
+      table foo
+      email foo@bar.example
+      private_key_path /path/to/key
+      project yourproject_id
+      dataset yourdataset_id
+
+      time_format %s
+      time_field  time
+
+      schema_path #{schema_path}
+      field_integer time
+      prevent_duplicate_load true
+
+      buffer_type memory
+    CONFIG
+    schema_fields = MultiJson.load(File.read(schema_path)).map(&:deep_symbolize_keys).tap do |h|
+      h[0][:type] = "INTEGER"
+      h[0][:mode] = "NULLABLE"
+    end
+
+    chunk = Fluent::MemoryBufferChunk.new("my.tag")
+    io = StringIO.new("hello")
+    mock(driver.instance).create_upload_source(chunk).yields(io)
+    mock.proxy(driver.instance).create_job_id(duck_type(:unique_id), "yourdataset_id", "foo", driver.instance.instance_variable_get(:@fields).to_a, 0, false)
+    writer = stub_writer(driver)
+    mock(writer).wait_load_job("yourproject_id", "yourdataset_id", "dummy_job_id", "foo") { nil }
+    mock(writer.client).insert_job('yourproject_id', {
+      configuration: {
+        load: {
+          destination_table: {
+            project_id: 'yourproject_id',
+            dataset_id: 'yourdataset_id',
+            table_id: 'foo',
+          },
+          schema: {
+            fields: schema_fields,
+          },
+          write_disposition: "WRITE_APPEND",
+          source_format: "NEWLINE_DELIMITED_JSON",
+          ignore_unknown_values: false,
+          max_bad_records: 0,
+        },
+      },
+      job_reference: {project_id: 'yourproject_id', job_id: satisfy { |x| x =~ /fluentd_job_.*/}} ,
+    }, {upload_source: io, content_type: "application/octet-stream", options: {timeout_sec: nil, open_timeout_sec: 60}}) do
+      s = stub!
+      job_reference_stub = stub!
+      s.job_reference { job_reference_stub }
+      job_reference_stub.job_id { "dummy_job_id" }
+      s
+    end
+
+    entry.each do |e|
+      chunk << MultiJson.dump(e) + "\n"
+    end
+
+    driver.instance.start
+    driver.instance.write(chunk)
+    driver.instance.shutdown
+  end
+
+  def test_write_for_load_with_retryable_error
+    schema_path = File.join(File.dirname(__FILE__), "testdata", "sudo.schema")
+    entry = {a: "b"}, {b: "c"}
+    driver = create_driver(<<-CONFIG)
+      method load
+      table foo
+      email foo@bar.example
+      private_key_path /path/to/key
+      project yourproject_id
+      dataset yourdataset_id
+
+      time_format %s
+      time_field  time
+
+      schema_path #{schema_path}
+      field_integer time
+
+      buffer_type memory
+    CONFIG
+    schema_fields = MultiJson.load(File.read(schema_path)).map(&:deep_symbolize_keys).tap do |h|
+      h[0][:type] = "INTEGER"
+      h[0][:mode] = "NULLABLE"
+    end
+
+    chunk = Fluent::MemoryBufferChunk.new("my.tag")
+    io = StringIO.new("hello")
+    mock(driver.instance).create_upload_source(chunk).yields(io)
+    writer = stub_writer(driver)
+    mock(writer.client).insert_job('yourproject_id', {
+      configuration: {
+        load: {
+          destination_table: {
+            project_id: 'yourproject_id',
+            dataset_id: 'yourdataset_id',
+            table_id: 'foo',
+          },
+          schema: {
+            fields: schema_fields,
+          },
+          write_disposition: "WRITE_APPEND",
+          source_format: "NEWLINE_DELIMITED_JSON",
+          ignore_unknown_values: false,
+          max_bad_records: 0,
+        }
+      }
+    }, {upload_source: io, content_type: "application/octet-stream", options: {timeout_sec: nil, open_timeout_sec: 60}}) do
+      s = stub!
+      job_reference_stub = stub!
+      s.job_reference { job_reference_stub }
+      job_reference_stub.job_id { "dummy_job_id" }
+      s
+    end
+
+    mock(writer.client).get_job('yourproject_id', 'dummy_job_id') do
+      s = stub!
+      status_stub = stub!
+      error_result = stub!
+
+      s.status { status_stub }
+      status_stub.state { "DONE" }
+      status_stub.error_result { error_result }
+      status_stub.errors { nil }
+      error_result.message { "error" }
+      error_result.reason { "backendError" }
+      s
+    end
+
+    entry.each do |e|
+      chunk << MultiJson.dump(e) + "\n"
+    end
+
+    driver.instance.start
+    assert_raise Fluent::BigQuery::Writer::RetryableError do
+      driver.instance.write(chunk)
+    end
+    driver.instance.shutdown
+  end
+
+  def test_write_for_load_with_not_retryable_error
+    schema_path = File.join(File.dirname(__FILE__), "testdata", "sudo.schema")
+    entry = {a: "b"}, {b: "c"}
+    driver = create_driver(<<-CONFIG)
+      method load
+      table foo
+      email foo@bar.example
+      private_key_path /path/to/key
+      project yourproject_id
+      dataset yourdataset_id
+
+      time_format %s
+      time_field  time
+
+      schema_path #{schema_path}
+      field_integer time
+
+      buffer_type memory
+      <secondary>
+        type file
+        path error
+        utc
+      </secondary>
+    CONFIG
+    schema_fields = MultiJson.load(File.read(schema_path)).map(&:deep_symbolize_keys).tap do |h|
+      h[0][:type] = "INTEGER"
+      h[0][:mode] = "NULLABLE"
+    end
+
+    chunk = Fluent::MemoryBufferChunk.new("my.tag")
+    io = StringIO.new("hello")
+    mock(driver.instance).create_upload_source(chunk).yields(io)
+    writer = stub_writer(driver)
+    mock(writer.client).insert_job('yourproject_id', {
+      configuration: {
+        load: {
+          destination_table: {
+            project_id: 'yourproject_id',
+            dataset_id: 'yourdataset_id',
+            table_id: 'foo',
+          },
+          schema: {
+            fields: schema_fields,
+          },
+          write_disposition: "WRITE_APPEND",
+          source_format: "NEWLINE_DELIMITED_JSON",
+          ignore_unknown_values: false,
+          max_bad_records: 0,
         }
-      }, {upload_source: io, content_type: "application/octet-stream"}) {
-        s = stub!
-        status_stub = stub!
-        s.status { status_stub }
-        status_stub.state { "DONE" }
-        status_stub.error_result { nil }
-        s
       }
+    }, {upload_source: io, content_type: "application/octet-stream", options: {timeout_sec: nil, open_timeout_sec: 60}}) do
+      s = stub!
+      job_reference_stub = stub!
+      s.job_reference { job_reference_stub }
+      job_reference_stub.job_id { "dummy_job_id" }
+      s
     end
 
+    mock(writer.client).get_job('yourproject_id', 'dummy_job_id') do
+      s = stub!
+      status_stub = stub!
+      error_result = stub!
+
+      s.status { status_stub }
+      status_stub.state { "DONE" }
+      status_stub.error_result { error_result }
+      status_stub.errors { nil }
+      error_result.message { "error" }
+      error_result.reason { "invalid" }
+      s
+    end
+
+    mock(driver.instance).flush_secondary(is_a(Fluent::Output))
+
     entry.each do |e|
       chunk << MultiJson.dump(e) + "\n"
     end
@@ -966,15 +1300,19 @@ class BigQueryOutputTest < Test::Unit::TestCase
       field_float   requesttime
       field_boolean bot_access,loginsession
     CONFIG
-    mock_client(driver) do |expect|
-      expect.insert_all_table_data('yourproject_id', 'yourdataset_id', 'foo_2014_08_20', {
-        rows: [entry[0]]
-      }, {})
-
-      expect.insert_all_table_data('yourproject_id', 'yourdataset_id', 'foo_2014_08_21', {
-        rows: [entry[1]]
-      }, {})
-    end
+
+    writer = stub_writer(driver)
+    mock(writer.client).insert_all_table_data('yourproject_id', 'yourdataset_id', 'foo_2014_08_20', {
+      rows: [entry[0]],
+      skip_invalid_rows: false,
+      ignore_unknown_values: false
+    }, {options: {timeout_sec: nil, open_timeout_sec: 60}}) { stub!.insert_errors { nil } }
+
+    mock(writer.client).insert_all_table_data('yourproject_id', 'yourdataset_id', 'foo_2014_08_21', {
+      rows: [entry[1]],
+      skip_invalid_rows: false,
+      ignore_unknown_values: false
+    }, {options: {timeout_sec: nil, open_timeout_sec: 60}}) { stub!.insert_errors { nil } }
 
     chunk = Fluent::MemoryBufferChunk.new("my.tag")
     entry.each do |object|
@@ -1085,23 +1423,13 @@ class BigQueryOutputTest < Test::Unit::TestCase
       auto_create_table true
       schema_path #{File.join(File.dirname(__FILE__), "testdata", "apache.schema")}
     CONFIG
-    mock_client(driver) do |expect|
-      expect.insert_all_table_data('yourproject_id', 'yourdataset_id', 'foo', {
-        rows: [message]
-      }, {}) {
-        raise Google::Apis::ServerError.new("Not found: Table yourproject_id:yourdataset_id.foo", status_code: 404, body: "Not found: Table yourproject_id:yourdataset_id.foo")
-      }
-      expect.insert_table('yourproject_id', 'yourdataset_id', {
-        table_reference: {
-          table_id: 'foo',
-        },
-        schema: {
-          fields: JSON.parse(File.read(File.join(File.dirname(__FILE__), "testdata", "apache.schema"))).map(&:deep_symbolize_keys),
-        }
-      }, {}) {
-        stub!
-      }
-    end
+    writer = stub_writer(driver)
+    mock(writer).insert_rows('yourproject_id', 'yourdataset_id', 'foo', [message], hash_including(
+      skip_invalid_rows: false,
+      ignore_unknown_values: false,
+    )) { raise Fluent::BigQuery::Writer::RetryableError.new(nil, Google::Apis::ServerError.new("Not found: Table yourproject_id:yourdataset_id.foo", status_code: 404, body: "Not found: Table yourproject_id:yourdataset_id.foo")) }
+    mock(writer).create_table('yourproject_id', 'yourdataset_id', 'foo', driver.instance.instance_variable_get(:@fields))
+
     chunk = Fluent::MemoryBufferChunk.new("my.tag")
     chunk << message.to_msgpack