/*
 * Copyright 2017 Google LLC
 *
 *  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 * in compliance with the License. You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software distributed under the
 * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 * express or implied. See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.google.cloud.hadoop.io.bigquery.mapred;

import static org.junit.Assert.assertThrows;
import static org.mockito.Matchers.any;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;

import com.google.gson.JsonObject;
import java.io.IOException;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.mapred.RecordWriter;
import org.apache.hadoop.mapred.Reporter;
import org.apache.hadoop.mapred.TaskAttemptContext;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;

/**
 * Unit tests for {@link BigQueryMapredRecordWriter}.
 */
@RunWith(JUnit4.class)
public class BigQueryMapredRecordWriterTest {

  @Mock private org.apache.hadoop.mapreduce.RecordWriter<
      LongWritable, JsonObject> mockRecordWriter;
  @Mock private TaskAttemptContext mockTaskAttemptContext;
  @Before public void setUp() {
    MockitoAnnotations.initMocks(this);
  }

  @After public void tearDown() {
    verifyNoMoreInteractions(mockRecordWriter);
    verifyNoMoreInteractions(mockTaskAttemptContext);
  }

  @Test public void testClose() throws IOException, InterruptedException {
    RecordWriter<LongWritable, JsonObject> recordWriter =
        new BigQueryMapredRecordWriter<LongWritable, JsonObject>(
        mockRecordWriter, mockTaskAttemptContext);
    Reporter reporter = null;   // unused by code under test

    recordWriter.close(reporter);
    verify(mockRecordWriter).close(any(TaskAttemptContext.class));

    doThrow(new IOException("test")).when(mockRecordWriter).close(any(TaskAttemptContext.class));
    assertThrows(IOException.class, () -> recordWriter.close(reporter));

    verify(mockRecordWriter, times(2)).close(any(TaskAttemptContext.class));
  }

  @Test public void testWrite() throws IOException, InterruptedException {
    RecordWriter<LongWritable, JsonObject> recordWriter =
        new BigQueryMapredRecordWriter<LongWritable, JsonObject>(
            mockRecordWriter, mockTaskAttemptContext);
    LongWritable key = new LongWritable(123);
    JsonObject value = new JsonObject();

    recordWriter.write(key, value);
    verify(mockRecordWriter).write(
        any(LongWritable.class), any(JsonObject.class));

    recordWriter.write(key, null);
    verify(mockRecordWriter, times(2)).write(
        any(LongWritable.class), any(JsonObject.class));

    doThrow(new IOException("test"))
        .when(mockRecordWriter)
        .write(any(LongWritable.class), any(JsonObject.class));
    assertThrows(IOException.class, () -> recordWriter.write(key, value));

    verify(mockRecordWriter, times(3)).write(any(LongWritable.class), any(JsonObject.class));
  }
}

@@ -682,16 +682,63 @@
 Throws;%0A
+import static org.mockito.ArgumentMatchers.eq;%0A
 import s
@@ -1407,19 +1407,16 @@
 ns;%0A%0A/**
-%0A *
  Unit te
@@ -1458,17 +1458,16 @@
 Writer%7D.
-%0A
  */%0A@Run
@@ -1589,23 +1589,16 @@
 dWriter%3C
-%0A      
 LongWrit
@@ -1692,16 +1692,17 @@
 ontext;%0A
+%0A
   @Befor
@@ -1702,16 +1702,18 @@
  @Before
+%0A 
  public 
@@ -1780,16 +1780,18 @@
   @After
+%0A 
  public 
@@ -1914,32 +1914,34 @@
 t);%0A  %7D%0A%0A  @Test
+%0A 
  public void tes
@@ -2116,32 +2116,36 @@
 bject%3E(%0A        
+    
 mockRecordWriter
@@ -2200,18 +2200,16 @@
  = null;
-  
  // unus
@@ -2594,16 +2594,18 @@
 %0A  @Test
+%0A 
  public 
@@ -2967,16 +2967,17 @@
 value);%0A
+%0A
     veri
@@ -3007,63 +3007,343 @@
 ite(
-%0A        any(LongWritable.class), any(JsonObject.class)
+eq(key), eq(value));%0A  %7D%0A%0A  @Test%0A  public void testWrite_nullValue() throws IOException, InterruptedException %7B%0A    RecordWriter%3CLongWritable, JsonObject%3E recordWriter =%0A        new BigQueryMapredRecordWriter%3CLongWritable, JsonObject%3E(%0A            mockRecordWriter, mockTaskAttemptContext);%0A    LongWritable key = new LongWritable(123
 );%0A%0A
@@ -3377,16 +3377,17 @@
  null);%0A
+%0A
     veri
@@ -3409,81 +3409,396 @@
 iter
-, times(2)).write(%0A        any(LongWritable.class), any(JsonObject.class)
+).write(eq(key), eq(null));%0A  %7D%0A%0A  @Test%0A  public void testWrite_throwException() throws IOException, InterruptedException %7B%0A    RecordWriter%3CLongWritable, JsonObject%3E recordWriter =%0A        new BigQueryMapredRecordWriter%3CLongWritable, JsonObject%3E(%0A            mockRecordWriter, mockTaskAttemptContext);%0A    LongWritable key = new LongWritable(123);%0A    JsonObject value = new JsonObject(
 );%0A%0A
@@ -4037,71 +4037,33 @@
 iter
-, times(3)).write(any(LongWritable.class), any(JsonObject.class
+).write(eq(key), eq(value
 ));%0A

Fix BigQuery tests after migration to Mockito 2