JsrPartitionHandlerTests.java example

Explorer
spring-batch-master
/*
 * Copyright 2013-2014 the original author or authors.
 *
 * 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 org.springframework.batch.core.jsr.partition;

import org.junit.Before;
import org.junit.Test;
import org.springframework.batch.core.ExitStatus;
import org.springframework.batch.core.JobExecution;
import org.springframework.batch.core.JobInstance;
import org.springframework.batch.core.JobInterruptedException;
import org.springframework.batch.core.StepExecution;
import org.springframework.batch.core.jsr.AbstractJsrTestCase;
import org.springframework.batch.core.jsr.configuration.support.BatchPropertyContext;
import org.springframework.batch.core.jsr.step.batchlet.BatchletSupport;
import org.springframework.batch.core.repository.JobRepository;
import org.springframework.batch.core.repository.support.MapJobRepositoryFactoryBean;
import org.springframework.batch.core.step.JobRepositorySupport;
import org.springframework.batch.core.step.StepSupport;

import javax.batch.api.BatchProperty;
import javax.batch.api.partition.PartitionAnalyzer;
import javax.batch.api.partition.PartitionCollector;
import javax.batch.api.partition.PartitionMapper;
import javax.batch.api.partition.PartitionPlan;
import javax.batch.api.partition.PartitionPlanImpl;
import javax.batch.api.partition.PartitionReducer;
import javax.batch.runtime.BatchStatus;
import javax.inject.Inject;
import java.io.Serializable;
import java.util.Collection;
import java.util.Properties;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.fail;

public class JsrPartitionHandlerTests extends AbstractJsrTestCase {

	private JsrPartitionHandler handler;
	private JobRepository repository = new JobRepositorySupport();
	private StepExecution stepExecution;
	private int count;
	private BatchPropertyContext propertyContext;
	private JsrStepExecutionSplitter stepSplitter;

	@Before
	public void setUp() throws Exception {
		JobExecution jobExecution = new JobExecution(1L);
		jobExecution.setJobInstance(new JobInstance(1L, "job"));
		stepExecution = new StepExecution("step1", jobExecution);
		stepSplitter = new JsrStepExecutionSplitter(repository, false, "step1", true);
		Analyzer.collectorData = "";
		Analyzer.status = "";
		count = 0;
		handler = new JsrPartitionHandler();
		handler.setStep(new StepSupport() {
			@Override
			public void execute(StepExecution stepExecution) throws JobInterruptedException {
				count++;
				stepExecution.setStatus(org.springframework.batch.core.BatchStatus.COMPLETED);
				stepExecution.setExitStatus(new ExitStatus("done"));
			}
		});
		propertyContext = new BatchPropertyContext();
		handler.setPropertyContext(propertyContext);
		repository = new MapJobRepositoryFactoryBean().getObject();
		handler.setJobRepository(repository);
		MyPartitionReducer.reset();
		CountingPartitionCollector.reset();
	}

	@Test
	public void testAfterPropertiesSet() throws Exception {
		handler = new JsrPartitionHandler();

		try {
			handler.afterPropertiesSet();
			fail("PropertyContext was not checked for");
		} catch(IllegalArgumentException iae) {
			assertEquals("A BatchPropertyContext is required", iae.getMessage());
		}

		handler.setPropertyContext(new BatchPropertyContext());

		try {
			handler.afterPropertiesSet();
			fail("Threads or mapper was not checked for");
		} catch(IllegalArgumentException iae) {
			assertEquals("Either a mapper implementation or the number of partitions/threads is required", iae.getMessage());
		}

		handler.setThreads(3);

		try {
			handler.afterPropertiesSet();
			fail("JobRepository was not checked for");
		} catch(IllegalArgumentException iae) {
			assertEquals("A JobRepository is required", iae.getMessage());
		}

		handler.setJobRepository(repository);
		handler.afterPropertiesSet();
	}

	@Test
	public void testHardcodedNumberOfPartitions() throws Exception {
		handler.setThreads(3);
		handler.setPartitions(3);
		handler.afterPropertiesSet();

		Collection<StepExecution> executions = handler.handle(stepSplitter, stepExecution);

		assertEquals(3, executions.size());
		assertEquals(3, count);
	}

	@Test
	public void testMapperProvidesPartitions() throws Exception {
		handler.setPartitionMapper(new PartitionMapper() {

			@Override
			public PartitionPlan mapPartitions() throws Exception {
				PartitionPlan plan = new PartitionPlanImpl();
				plan.setPartitions(3);
				plan.setThreads(0);
				return plan;
			}
		});

		handler.afterPropertiesSet();

		Collection<StepExecution> executions = handler.handle(new JsrStepExecutionSplitter(repository, false, "step1", true), stepExecution);

		assertEquals(3, executions.size());
		assertEquals(3, count);
	}

	@Test
	public void testMapperProvidesPartitionsAndThreads() throws Exception {
		handler.setPartitionMapper(new PartitionMapper() {

			@Override
			public PartitionPlan mapPartitions() throws Exception {
				PartitionPlan plan = new PartitionPlanImpl();
				plan.setPartitions(3);
				plan.setThreads(1);
				return plan;
			}
		});

		handler.afterPropertiesSet();

		Collection<StepExecution> executions = handler.handle(new JsrStepExecutionSplitter(repository, false, "step1", true), stepExecution);

		assertEquals(3, executions.size());
		assertEquals(3, count);
	}

	@Test
	public void testMapperWithProperties() throws Exception {
		handler.setPartitionMapper(new PartitionMapper() {

			@Override
			public PartitionPlan mapPartitions() throws Exception {
				PartitionPlan plan = new PartitionPlanImpl();
				Properties [] props = new Properties[2];
				props[0] = new Properties();
				props[0].put("key1", "value1");
				props[1] = new Properties();
				props[1].put("key1", "value2");
				plan.setPartitionProperties(props);
				plan.setPartitions(3);
				plan.setThreads(1);
				return plan;
			}
		});

		handler.afterPropertiesSet();

		Collection<StepExecution> executions = handler.handle(new JsrStepExecutionSplitter(repository, false, "step1", true), stepExecution);

		assertEquals(3, executions.size());
		assertEquals(3, count);
		assertEquals("value1", propertyContext.getStepProperties("step1:partition0").get("key1"));
		assertEquals("value2", propertyContext.getStepProperties("step1:partition1").get("key1"));
	}

	@Test
	public void testAnalyzer() throws Exception {
		Queue<Serializable> queue = new ConcurrentLinkedQueue<Serializable>();
		queue.add("foo");
		queue.add("bar");

		handler.setPartitionDataQueue(queue);
		handler.setThreads(2);
		handler.setPartitions(2);
		handler.setPartitionAnalyzer(new Analyzer());
		handler.afterPropertiesSet();

		Collection<StepExecution> executions = handler.handle(new JsrStepExecutionSplitter(repository, false, "step1", true), stepExecution);

		assertEquals(2, executions.size());
		assertEquals(2, count);
		assertEquals("foobar", Analyzer.collectorData);
		assertEquals("COMPLETEDdone", Analyzer.status);
	}

	@Test
	public void testRestartNoOverride() throws Exception {
		javax.batch.runtime.JobExecution execution1 = runJob("jsrPartitionHandlerRestartWithOverrideJob", null, 1000000L);
		assertEquals(BatchStatus.FAILED, execution1.getBatchStatus());
		assertEquals(1, MyPartitionReducer.beginCount);
		assertEquals(0, MyPartitionReducer.beforeCount);
		assertEquals(1, MyPartitionReducer.rollbackCount);
		assertEquals(1, MyPartitionReducer.afterCount);
		assertEquals(3, CountingPartitionCollector.collected);

		MyPartitionReducer.reset();
		CountingPartitionCollector.reset();

		javax.batch.runtime.JobExecution execution2 = restartJob(execution1.getExecutionId(), null, 1000000L);
		assertEquals(BatchStatus.COMPLETED, execution2.getBatchStatus());
		assertEquals(1, MyPartitionReducer.beginCount);
		assertEquals(1, MyPartitionReducer.beforeCount);
		assertEquals(0, MyPartitionReducer.rollbackCount);
		assertEquals(1, MyPartitionReducer.afterCount);
		assertEquals(1, CountingPartitionCollector.collected);
	}


	@Test
	public void testRestartOverride() throws Exception {
		Properties jobParameters = new Properties();
		jobParameters.put("mapper.override", "true");

		javax.batch.runtime.JobExecution execution1 = runJob("jsrPartitionHandlerRestartWithOverrideJob", jobParameters, 1000000L);
		assertEquals(BatchStatus.FAILED, execution1.getBatchStatus());
		assertEquals(1, MyPartitionReducer.beginCount);
		assertEquals(0, MyPartitionReducer.beforeCount);
		assertEquals(1, MyPartitionReducer.rollbackCount);
		assertEquals(1, MyPartitionReducer.afterCount);
		assertEquals(3, CountingPartitionCollector.collected);

		MyPartitionReducer.reset();
		CountingPartitionCollector.reset();

		javax.batch.runtime.JobExecution execution2 = restartJob(execution1.getExecutionId(), jobParameters, 1000000L);
		assertEquals(BatchStatus.COMPLETED, execution2.getBatchStatus());
		assertEquals(1, MyPartitionReducer.beginCount);
		assertEquals(1, MyPartitionReducer.beforeCount);
		assertEquals(0, MyPartitionReducer.rollbackCount);
		assertEquals(1, MyPartitionReducer.afterCount);
		assertEquals(5, CountingPartitionCollector.collected);
	}

	public static class CountingPartitionCollector implements PartitionCollector {

		public static int collected = 0;

		public static void reset() {
			collected = 0;
		}

		@Override
		public Serializable collectPartitionData() throws Exception {
			collected++;

			return null;
		}
	}

	public static class MyPartitionReducer implements PartitionReducer {

		public static int beginCount = 0;
		public static int beforeCount = 0;
		public static int rollbackCount = 0;
		public static int afterCount = 0;

		public static void reset() {
			beginCount = 0;
			beforeCount = 0;
			rollbackCount = 0;
			afterCount = 0;
		}

		@Override
		public void beginPartitionedStep() throws Exception {
			beginCount++;
		}

		@Override
		public void beforePartitionedStepCompletion() throws Exception {
			beforeCount++;
		}

		@Override
		public void rollbackPartitionedStep() throws Exception {
			rollbackCount++;
		}

		@Override
		public void afterPartitionedStepCompletion(PartitionStatus status)
				throws Exception {
			afterCount++;
		}
	}

	public static class MyPartitionMapper implements PartitionMapper {

		private static int count = 0;

		@Inject
		@BatchProperty
		String overrideString = "false";

		@Override
		public PartitionPlan mapPartitions() throws Exception {
			count++;

			PartitionPlan plan = new PartitionPlanImpl();

			if(count % 2 == 1) {
				plan.setPartitions(3);
				plan.setThreads(3);
			} else {
				plan.setPartitions(5);
				plan.setThreads(5);
			}

			plan.setPartitionsOverride(Boolean.valueOf(overrideString));

			Properties[] props = new Properties[3];
			props[0] = new Properties();
			props[1] = new Properties();
			props[2] = new Properties();

			if(count % 2 == 1) {
				props[1].put("fail", "true");
			}

			plan.setPartitionProperties(props);
			return plan;
		}
	}

	public static class MyBatchlet extends BatchletSupport {
		@Inject
		@BatchProperty
		String fail;

		@Override
		public String process() {
			if("true".equalsIgnoreCase(fail)) {
				throw new RuntimeException("Expected");
			}

			return null;
		}
	}

	public static class Analyzer implements PartitionAnalyzer {

		public static String collectorData;
		public static String status;

		@Override
		public void analyzeCollectorData(Serializable data) throws Exception {
			collectorData = collectorData + data;
		}

		@Override
		public void analyzeStatus(BatchStatus batchStatus, String exitStatus)
				throws Exception {
			status = batchStatus + exitStatus;
		}
	}
}