DatabaseOperatorImplTest.java

/*
 * Copyright 2017 LinkedIn Corp.
 *
 * 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
 *
 * http://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 azkaban.db;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import javax.sql.DataSource;
import org.apache.commons.dbutils.QueryRunner;
import org.apache.commons.dbutils.ResultSetHandler;
import org.apache.commons.dbutils.handlers.ScalarHandler;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import static org.mockito.Mockito.*;


public class DatabaseOperatorImplTest {

  private AzkabanDataSource datasource = new AzDBTestUtility.EmbeddedH2BasicDataSource();

  private DatabaseOperator dbOperator;
  private QueryRunner queryRunner;
  private Connection conn;

  private ResultSetHandler<Integer> handler = rs -> {
    if (!rs.next()) {
      return 0;
    }
    return rs.getInt(1);
  };

  private static final List<Integer> list = new ArrayList<>();

  private static int index_1 = 3;
  private static int index_2 = 15;

  @Before
  public void setUp() throws Exception {
    queryRunner = mock(QueryRunner.class);

    conn = datasource.getConnection();
    DataSource mockDataSource = mock(datasource.getClass());

    when(queryRunner.getDataSource()).thenReturn(mockDataSource);
    when(mockDataSource.getConnection()).thenReturn(conn);

    this.dbOperator = new DatabaseOperatorImpl(queryRunner);

    list.add(index_1);
    list.add(index_2);

    // valid query returns correct value
    when(queryRunner.query("select * from blah where ? = ?", handler, "id", 2)).thenReturn(index_2);

    // If select an non-existing entry, handler returns 0.
    when(queryRunner.query("select * from blah where ? = ?", handler, "id", 3)).thenReturn(0);

    //If typos, throw Exceptions.
    doThrow(SQLException.class).when(queryRunner).query("sele * from blah where ? = ?", handler, "id", 2);

    doAnswer(invocation -> {
      index_1 = 26;
      return 1;
    }).when(queryRunner).update("update blah set ? = ?", "1", 26);
  }

  @Test
  public void testValidQuery() throws Exception {
    int res = dbOperator.query("select * from blah where ? = ?", handler, "id", 2);
    Assert.assertEquals(15, res);
    verify(queryRunner).query("select * from blah where ? = ?", handler, "id", 2);
  }

  @Test
  public void testInvalidQuery() throws Exception {
    int res = dbOperator.query("select * from blah where ? = ?", handler, "id", 3);
    Assert.assertEquals(0, res);
  }

  @Test(expected = SQLException.class)
  public void testTypoSqlStatement() throws Exception {
    System.out.println("testTypoSqlStatement");
    dbOperator.query("sele * from blah where ? = ?", handler, "id", 2);
  }

  @Test
  public void testTransaction() throws Exception {
    when(queryRunner.update(conn, "update blah set ? = ?", "1", 26)).thenReturn(1);
    when(queryRunner.query(conn, "select * from blah where ? = ?", handler, "id", 1)).thenReturn(26);

    SQLTransaction<Integer> transaction = transOperator -> {
      transOperator.update("update blah set ? = ?", "1", 26);
      return transOperator.query("select * from blah where ? = ?", handler, "id", 1);
    };

    int res = dbOperator.transaction(transaction);
    Assert.assertEquals(26, res);
  }

  @Test
  public void testValidUpdate() throws Exception {
    int res = dbOperator.update("update blah set ? = ?", "1", 26);

    // 1 row is affected
    Assert.assertEquals(1, res);
    Assert.assertEquals(26, index_1);
    verify(queryRunner).update("update blah set ? = ?", "1", 26);
  }

  @Test
  public void testInvalidUpdate() throws Exception {
    int res = dbOperator.update("update blah set ? = ?", "3", 26);

    // 0 row is affected
    Assert.assertEquals(0, res);
  }
}