azkaban.db.DatabaseOperatorImplTest.java Source code

Introduction

Here is the source code for azkaban.db.DatabaseOperatorImplTest.java

Source

/*
 * 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);
    }
}