BigQueryPullServiceImplTest.java

/*
 * This file is a component of thundr, a software library from 3wks.
 * Read more: http://www.3wks.com.au/thundr
 * Copyright (C) 2013 3wks, <thundr@3wks.com.au>
 *
 * 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 com.threewks.thundr.bigquery;

import static org.hamcrest.Matchers.is;
import static org.junit.Assert.assertThat;
import static org.mockito.Mockito.*;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import org.junit.Before;
import org.junit.Test;
import org.mockito.Mockito;

import com.google.api.services.bigquery.Bigquery;
import com.google.api.services.bigquery.Bigquery.Jobs;
import com.google.api.services.bigquery.Bigquery.Jobs.Get;
import com.google.api.services.bigquery.Bigquery.Jobs.GetQueryResults;
import com.google.api.services.bigquery.Bigquery.Jobs.Insert;
import com.google.api.services.bigquery.Bigquery.Jobs.Query;
import com.google.api.services.bigquery.model.Job;
import com.google.api.services.bigquery.model.JobReference;
import com.google.api.services.bigquery.model.JobStatus;
import com.google.api.services.bigquery.model.QueryRequest;
import com.google.api.services.bigquery.model.QueryResponse;
import com.google.api.services.bigquery.model.TableCell;
import com.google.api.services.bigquery.model.TableFieldSchema;
import com.google.api.services.bigquery.model.TableRow;
import com.google.api.services.bigquery.model.TableSchema;

public class BigQueryPullServiceImplTest {

	BigQueryPullService service;
	private String bigQueryProjectId = "big-query-project";

	@Before
	public void before() throws IOException {
		Bigquery bigQuery = mock(Bigquery.class);
		Jobs jobs = mock(Jobs.class);
		Query query = mock(Query.class);

		when(jobs.query(Mockito.anyString(), Mockito.any(QueryRequest.class))).thenReturn(query);
		when(bigQuery.jobs()).thenReturn(jobs);
		QueryResponse queryResponse = new QueryResponse();
		queryResponse.setJobComplete(true);
		when(query.execute()).thenReturn(queryResponse);

		TableSchema schema = createSchema("field1", "field2", "field3");
		queryResponse.setSchema(schema);

		List<TableRow> rows = new ArrayList<TableRow>();
		rows.add(createRow("abc", "def", "ghi"));
		rows.add(createRow("jkl", "mno", "pqr"));
		queryResponse.setRows(rows);

		Insert insert = mock(Insert.class);
		JobReference jobReference = new JobReference();
		jobReference.setJobId("job_123");
		Job job = new Job();
		JobStatus status = new JobStatus();
		status.setState("DONE");
		job.setStatus(status);
		job.setJobReference(jobReference);
		when(insert.execute()).thenReturn(job);
		when(jobs.insert(Mockito.anyString(), Mockito.any(Job.class))).thenReturn(insert);

		Get get = mock(Get.class);
		when(get.execute()).thenReturn(job);
		when(jobs.get(bigQueryProjectId, "job_123")).thenReturn(get);

		GetQueryResults results = mock(GetQueryResults.class);
		when(jobs.getQueryResults(bigQueryProjectId, "job_123")).thenReturn(results);

		service = new BigQueryPullServiceImpl(bigQuery, bigQueryProjectId);
	}

	@Test
	public void shouldExecuteQuery() throws Exception {
		List<List<String>> results = service.query("select * from table", true);
		assertThat(results.size(), is(3));
		assertRow(results.get(0), "field1", "field2", "field3");
		assertRow(results.get(1), "abc", "def", "ghi");
		assertRow(results.get(2), "jkl", "mno", "pqr");
	}

	@Test
	public void shouldExecuteQueryAsync() throws Exception {
		String jobId = service.queryAsync("select * from table");
		assertThat(jobId, is("job_123"));
	}

	@Test
	public void shouldGetJobStatus() throws Exception {
		String status = service.getJobStatus("job_123");
		assertThat(status, is("DONE"));
	}

	@Test
	public void shouldNotIncludeColumnNames() throws Exception {
		List<List<String>> results = service.query("select * from table", false);
		assertThat(results.size(), is(2));
		assertRow(results.get(0), "abc", "def", "ghi");
		assertRow(results.get(1), "jkl", "mno", "pqr");
	}

	private TableSchema createSchema(String... fieldNames) {
		TableSchema schema = new TableSchema();
		List<TableFieldSchema> fields = new ArrayList<TableFieldSchema>();
		for (String fieldName : fieldNames) {
			TableFieldSchema field = new TableFieldSchema();
			field.setName(fieldName);
			fields.add(field);
		}
		schema.setFields(fields);
		return schema;
	}

	private TableRow createRow(String... values) {
		TableRow row = new TableRow();
		List<TableCell> tableCells = new ArrayList<TableCell>();
		for (String value : values) {
			TableCell cell = new TableCell();
			cell.setV(value);
			tableCells.add(cell);
		}
		row.setF(tableCells);
		return row;
	}

	private void assertRow(List<String> row, String... values) {
		assertThat(row.size(), is(values.length));
		for (int i = 0; i < values.length; i++) {
			assertThat(row.get(i), is(values[i]));
		}
	}
}