BigQueryPushServiceImplTest.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.ByteArrayInputStream;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.TimeUnit;

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

import com.atomicleopard.expressive.Expressive;
import com.google.api.client.http.ByteArrayContent;
import com.google.api.services.bigquery.Bigquery;
import com.google.api.services.bigquery.Bigquery.Jobs;
import com.google.api.services.bigquery.Bigquery.Jobs.Insert;
import com.google.api.services.bigquery.model.Job;
import com.google.api.services.bigquery.model.JobConfigurationLoad;
import com.google.api.services.bigquery.model.TableReference;
import com.google.appengine.api.taskqueue.Queue;
import com.google.appengine.api.taskqueue.TaskHandle;
import com.google.appengine.api.taskqueue.TaskOptions;

public class BigQueryPushServiceImplTest {

	private BigQueryPushService bigQueryPushService;
	private Queue bigQueryQueue;
	private Bigquery bigQuery;
	private List<TaskHandle> tasks;
	private Insert insert;
	private Jobs jobs;

	private String bigQueryProjectId = "projectId";
	private String bigQueryDatasetId = "datasetId";

	private ArgumentCaptor<String> projectIdCaptor = ArgumentCaptor.forClass(String.class);
	private ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class);
	private ArgumentCaptor<ByteArrayContent> contentCaptor = ArgumentCaptor.forClass(ByteArrayContent.class);

	@Before
	public void before() throws Exception {

		bigQuery = mock(Bigquery.class);
		bigQueryQueue = mock(Queue.class);
		jobs = mock(Jobs.class);
		insert = mock(Insert.class);

		when(bigQuery.jobs()).thenReturn(jobs);
		when(jobs.insert(projectIdCaptor.capture(), jobCaptor.capture(), contentCaptor.capture())).thenReturn(insert);

		tasks = new ArrayList<TaskHandle>();
		tasks.add(createTestTask(Expressive.<String, String> map("tableId", "table1", "data", "a,b,c"), "task1", "bigQueryQueue"));
		tasks.add(createTestTask(Expressive.<String, String> map("tableId", "table1", "data", "d,e,f"), "task2", "bigQueryQueue"));
		tasks.add(createTestTask(Expressive.<String, String> map("tableId", "table2", "data", "1,2,3"), "task3", "bigQueryQueue"));
		when(bigQueryQueue.leaseTasks(10, TimeUnit.MINUTES, 1000)).thenReturn(tasks);

		bigQueryPushService = new BigQueryPushServiceImpl(bigQuery, bigQueryQueue, bigQueryProjectId, bigQueryDatasetId);
	}

	@Test
	public void shouldProcessTasks() throws Exception {
		int count = bigQueryPushService.processTasks();

		assertThat(count, is(3));

		// expect 2 insert() calls - one for table1 and another for table2
		verify(insert, times(2)).execute();

		List<String> projectIds = projectIdCaptor.getAllValues();
		List<Job> jobs = jobCaptor.getAllValues();
		List<ByteArrayContent> contents = contentCaptor.getAllValues();

		assertThat(projectIds.get(0), is("projectId"));
		assertThat(projectIds.get(1), is("projectId"));
		assertJob(jobs.get(0), "table1", "datasetId", "projectId");
		assertJob(jobs.get(1), "table2", "datasetId", "projectId");
		assertContent(contents.get(0), "a,b,c\nd,e,f\n");
		assertContent(contents.get(1), "1,2,3\n");
	}

	@Test
	@SuppressWarnings("unchecked")
	public void shouldRegisterConverters() {

		EventConverter<String> converter = mock(EventConverter.class);
		when(converter.getEventClass()).thenReturn(String.class);
		when(converter.getCsvData(Mockito.anyString())).thenReturn("1,2,3");
		when(converter.getTableId()).thenReturn("tableId");

		bigQueryPushService.registerEventConverter(converter);
		bigQueryPushService.trackEvent("Test Event");

		verify(converter).getEventClass();
		verify(converter).getCsvData("Test Event");
		verify(converter).getTableId();
	}

	/**
	 * Assert that a {@link ByteArrayContent} has the expected content and type.
	 * 
	 * @param content the {@link ByteArrayContent} to inspect.
	 * @param value the expected content value.
	 */
	private void assertContent(ByteArrayContent content, String value) throws Exception {
		assertThat(content.getType(), is("application/octet-stream"));
		ByteArrayInputStream inputStream = (ByteArrayInputStream) content.getInputStream();
		byte[] data = new byte[inputStream.available()];
		inputStream.read(data);
		assertThat(data, is(value.getBytes()));
	}

	/**
	 * Assert that a {@link Job} contains the expected data.
	 * 
	 * @param job the job to inspect.
	 * @param tableId the expected table id.
	 * @param datasetId the expected dataset id.
	 * @param projectId the expected project id.
	 */
	private void assertJob(Job job, String tableId, String datasetId, String projectId) {
		JobConfigurationLoad load = job.getConfiguration().getLoad();
		assertThat(load.getAllowQuotedNewlines(), is(true));
		TableReference destinationTable = load.getDestinationTable();
		assertThat(destinationTable.getTableId(), is(tableId));
		assertThat(destinationTable.getDatasetId(), is(datasetId));
		assertThat(destinationTable.getProjectId(), is(projectId));
	}

	private static TaskHandle createTestTask(Map<String, String> params, String taskName, String queueName) {
		TaskOptions options = TaskOptions.Builder.withDefaults();

		// build the payload using the parameter map
		StringBuilder payload = new StringBuilder();
		Iterator<Entry<String, String>> iterator = params.entrySet().iterator();
		while (iterator.hasNext()) {
			Entry<String, String> entry = iterator.next();
			payload.append(String.format("%s=%s&", entry.getKey(), entry.getValue()));
		}
		options.taskName(taskName);
		options.payload(payload.toString());
		TaskHandle task = new TaskHandle(options, queueName);
		return task;
	}

}