BATCH-220: Rationalise the exception classifiers so they can be used in a retry

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/xml/StaxEventItemWriterTests.java

@@ -11,7 +11,6 @@ import static org.junit.Assert.fail;
 import java.io.File;
 import java.io.IOException;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 
 import javax.xml.stream.XMLEventFactory;
@@ -48,7 +47,7 @@ public class StaxEventItemWriterTests {
 	// test item for writing to output
 	private Object item = new Object() {
 		public String toString() {
-			return ClassUtils.getShortName(StaxEventItemWriter.class)+"-testString";
+			return ClassUtils.getShortName(StaxEventItemWriter.class) + "-testString";
 		}
 	};
 
@@ -73,7 +72,7 @@ public class StaxEventItemWriterTests {
 		writer.write(items);
 		writer.close(executionContext);
 		String content = outputFileContent();
-		assertTrue("Wrong content: "+content, content.contains(TEST_STRING));
+		assertTrue("Wrong content: " + content, content.contains(TEST_STRING));
 	}
 
 	/**
@@ -96,7 +95,7 @@ public class StaxEventItemWriterTests {
 		// check the output is concatenation of 'before restart' and 'after
 		// restart' writes.
 		String outputFile = outputFileContent();
-		
+
 		assertEquals(2, StringUtils.countOccurrencesOf(outputFile, TEST_STRING));
 	}
 
@@ -107,13 +106,13 @@ public class StaxEventItemWriterTests {
 	public void testWriteWithHeader() throws Exception {
 		Object header1 = new Object();
 		Object header2 = new Object();
-		writer.setHeaderItems(new Object[] {header1, header2});
+		writer.setHeaderItems(new Object[] { header1, header2 });
 		writer.open(executionContext);
 		writer.write(items);
 		String content = outputFileContent();
-		assertTrue("Wrong content: "+content, content.contains(("<!--" + header1 + "-->")));
-		assertTrue("Wrong content: "+content, content.contains(("<!--" + header2 + "-->")));
-		assertTrue("Wrong content: "+content, content.contains(TEST_STRING));
+		assertTrue("Wrong content: " + content, content.contains(("<!--" + header1 + "-->")));
+		assertTrue("Wrong content: " + content, content.contains(("<!--" + header2 + "-->")));
+		assertTrue("Wrong content: " + content, content.contains(TEST_STRING));
 	}
 
 	/**
@@ -123,8 +122,7 @@ public class StaxEventItemWriterTests {
 	public void testStreamContext() throws Exception {
 		writer.open(executionContext);
 		final int NUMBER_OF_RECORDS = 10;
-		assertFalse(executionContext.containsKey(ClassUtils.getShortName(StaxEventItemWriter.class)
-				+ ".record.count"));
+		assertFalse(executionContext.containsKey(ClassUtils.getShortName(StaxEventItemWriter.class) + ".record.count"));
 		for (int i = 1; i <= NUMBER_OF_RECORDS; i++) {
 			writer.write(items);
 			writer.update(executionContext);
@@ -141,25 +139,21 @@ public class StaxEventItemWriterTests {
 	@Test
 	public void testOpenAndClose() throws Exception {
 		writer.setRootTagName("testroot");
-		writer.setRootElementAttributes(new HashMap<String, String>() {
-			{
-				put("attribute", "value");
-			}
-		});
+		writer.setRootElementAttributes(Collections.<String, String> singletonMap("attribute", "value"));
 		writer.open(executionContext);
 		writer.close(null);
 		String content = outputFileContent();
 		assertTrue(content.contains("<testroot attribute=\"value\">"));
 		assertTrue(content.endsWith("</testroot>"));
 	}
-	
+
 	@Test
 	public void testNonExistantResource() throws Exception {
 		Resource doesntExist = createMock(Resource.class);
 		expect(doesntExist.getFile()).andReturn(File.createTempFile("arbitrary", null));
 		expect(doesntExist.exists()).andReturn(false);
 		replay(doesntExist);
-		
+
 		writer.setResource(doesntExist);
 
 		try {

spring-batch-infrastructure/src/test/java/org/springframework/batch/repeat/exception/LogOrRethrowExceptionHandlerTests.java

@@ -20,34 +20,37 @@ import java.io.StringWriter;
 
 import junit.framework.TestCase;
 
-import org.apache.log4j.Level;
 import org.apache.log4j.Logger;
 import org.apache.log4j.SimpleLayout;
 import org.apache.log4j.WriterAppender;
 import org.springframework.batch.repeat.RepeatContext;
-import org.springframework.batch.support.ExceptionClassifierSupport;
+import org.springframework.batch.repeat.exception.LogOrRethrowExceptionHandler.Level;
+import org.springframework.batch.support.ClassifierSupport;
 
 public class LogOrRethrowExceptionHandlerTests extends TestCase {
 
 	private LogOrRethrowExceptionHandler handler = new LogOrRethrowExceptionHandler();
+
 	private StringWriter writer;
+
 	private RepeatContext context = null;
-	
+
 	protected void setUp() throws Exception {
 		super.setUp();
 		Logger logger = Logger.getLogger(LogOrRethrowExceptionHandler.class);
-		logger.setLevel(Level.DEBUG);
+		logger.setLevel(org.apache.log4j.Level.DEBUG);
 		writer = new StringWriter();
 		logger.removeAllAppenders();
 		logger.getParent().removeAllAppenders();
 		logger.addAppender(new WriterAppender(new SimpleLayout(), writer));
 	}
-	
+
 	public void testRuntimeException() throws Throwable {
 		try {
 			handler.handleException(context, new RuntimeException("Foo"));
 			fail("Expected RuntimeException");
-		} catch (RuntimeException e) {
+		}
+		catch (RuntimeException e) {
 			assertEquals("Foo", e.getMessage());
 		}
 	}
@@ -56,15 +59,16 @@ public class LogOrRethrowExceptionHandlerTests extends TestCase {
 		try {
 			handler.handleException(context, new Error("Foo"));
 			fail("Expected Error");
-		} catch (Error e) {
+		}
+		catch (Error e) {
 			assertEquals("Foo", e.getMessage());
 		}
 	}
-	
+
 	public void testNotRethrownErrorLevel() throws Throwable {
-		handler.setExceptionClassifier(new ExceptionClassifierSupport() {
-			public String classify(Throwable throwable) {
-				return LogOrRethrowExceptionHandler.ERROR;
+		handler.setExceptionClassifier(new ClassifierSupport<Throwable,Level>(Level.RETHROW) {
+			public Level classify(Throwable throwable) {
+				return Level.ERROR;
 			}
 		});
 		// No exception...
@@ -73,20 +77,9 @@ public class LogOrRethrowExceptionHandlerTests extends TestCase {
 	}
 
 	public void testNotRethrownWarnLevel() throws Throwable {
-		handler.setExceptionClassifier(new ExceptionClassifierSupport() {
-			public String classify(Throwable throwable) {
-				return LogOrRethrowExceptionHandler.WARN;
-			}
-		});
-		// No exception...
-		handler.handleException(context, new Error("Foo"));
-		assertNotNull(writer.toString());
-	}
-	
-	public void testNotRethrownDebugLevel() throws Throwable {
-		handler.setExceptionClassifier(new ExceptionClassifierSupport() {
-			public String classify(Throwable throwable) {
-				return LogOrRethrowExceptionHandler.DEBUG;
+		handler.setExceptionClassifier(new ClassifierSupport<Throwable,Level>(Level.RETHROW) {
+			public Level classify(Throwable throwable) {
+				return Level.WARN;
 			}
 		});
 		// No exception...
@@ -94,18 +87,15 @@ public class LogOrRethrowExceptionHandlerTests extends TestCase {
 		assertNotNull(writer.toString());
 	}
 
-	public void testUnclassifiedException() throws Throwable {
-		handler.setExceptionClassifier(new ExceptionClassifierSupport() {
-			public String classify(Throwable throwable) {
-				return "DEFAULT";
+	public void testNotRethrownDebugLevel() throws Throwable {
+		handler.setExceptionClassifier(new ClassifierSupport<Throwable,Level>(Level.RETHROW) {
+			public Level classify(Throwable throwable) {
+				return Level.DEBUG;
 			}
 		});
-		try {
-			handler.handleException(context, new Error("Foo"));
-			fail("Expected IllegalStateException");
-		} catch (IllegalStateException e) {
-			assertTrue(e.getMessage().toLowerCase().indexOf("unclassified")>=0);
-		}
+		// No exception...
+		handler.handleException(context, new Error("Foo"));
+		assertNotNull(writer.toString());
 	}
 
 }

spring-batch-infrastructure/src/test/java/org/springframework/batch/repeat/exception/RethrowOnThresholdExceptionHandlerTests.java

@@ -16,60 +16,60 @@
 
 package org.springframework.batch.repeat.exception;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.fail;
+
 import java.util.Collections;
+import java.util.concurrent.atomic.AtomicInteger;
 
-import junit.framework.TestCase;
-
+import org.junit.Test;
 import org.springframework.batch.repeat.RepeatContext;
-import org.springframework.batch.repeat.context.RepeatContextCounter;
 import org.springframework.batch.repeat.context.RepeatContextSupport;
-import org.springframework.batch.support.ExceptionClassifierSupport;
 
-public class RethrowOnThresholdExceptionHandlerTests extends TestCase {
+public class RethrowOnThresholdExceptionHandlerTests {
 
 	private RethrowOnThresholdExceptionHandler handler = new RethrowOnThresholdExceptionHandler();
+
 	private RepeatContext parent = new RepeatContextSupport(null);
+
 	private RepeatContext context = new RepeatContextSupport(parent);
-	
+
+	@Test
 	public void testRuntimeException() throws Throwable {
 		try {
 			handler.handleException(context, new RuntimeException("Foo"));
 			fail("Expected RuntimeException");
-		} catch (RuntimeException e) {
+		}
+		catch (RuntimeException e) {
 			assertEquals("Foo", e.getMessage());
 		}
 	}
 
+	@Test
 	public void testError() throws Throwable {
 		try {
 			handler.handleException(context, new Error("Foo"));
 			fail("Expected Error");
-		} catch (Error e) {
+		}
+		catch (Error e) {
 			assertEquals("Foo", e.getMessage());
 		}
 	}
-	
+
+	@Test
 	public void testNotRethrownWithThreshold() throws Throwable {
-		handler.setExceptionClassifier(new ExceptionClassifierSupport() {
-			public String classify(Throwable throwable) {
-				return "RuntimeException";
-			}
-		});
-		handler.setThresholds(Collections.singletonMap("RuntimeException", new Integer(1)));
+		handler.setThresholds(Collections.<Class<? extends Throwable>, Integer> singletonMap(Exception.class, 1));
 		// No exception...
 		handler.handleException(context, new RuntimeException("Foo"));
-		RepeatContextCounter counter = new RepeatContextCounter(context, RethrowOnThresholdExceptionHandler.class.getName() + ".RuntimeException");
+		AtomicInteger counter = (AtomicInteger) context.getAttribute(context.attributeNames()[0]);
 		assertNotNull(counter);
-		assertEquals(1, counter.getCount());
+		assertEquals(1, counter.get());
 	}
-	
+
+	@Test
 	public void testRethrowOnThreshold() throws Throwable {
-		handler.setExceptionClassifier(new ExceptionClassifierSupport() {
-			public String classify(Throwable throwable) {
-				return "RuntimeException";
-			}
-		});
-		handler.setThresholds(Collections.singletonMap("RuntimeException", new Integer(2)));
+		handler.setThresholds(Collections.<Class<? extends Throwable>, Integer> singletonMap(Exception.class, 2));
 		// No exception...
 		handler.handleException(context, new RuntimeException("Foo"));
 		handler.handleException(context, new RuntimeException("Foo"));
@@ -81,14 +81,10 @@ public class RethrowOnThresholdExceptionHandlerTests extends TestCase {
 			assertEquals("Foo", e.getMessage());
 		}
 	}
-	
+
+	@Test
 	public void testNotUseParent() throws Throwable {
-		handler.setExceptionClassifier(new ExceptionClassifierSupport() {
-			public String classify(Throwable throwable) {
-				return "RuntimeException";
-			}
-		});
-		handler.setThresholds(Collections.singletonMap("RuntimeException", new Integer(1)));
+		handler.setThresholds(Collections.<Class<? extends Throwable>, Integer> singletonMap(Exception.class, 1));
 		// No exception...
 		handler.handleException(context, new RuntimeException("Foo"));
 		context = new RepeatContextSupport(parent);
@@ -101,13 +97,9 @@ public class RethrowOnThresholdExceptionHandlerTests extends TestCase {
 		}
 	}
 
+	@Test
 	public void testUseParent() throws Throwable {
-		handler.setExceptionClassifier(new ExceptionClassifierSupport() {
-			public String classify(Throwable throwable) {
-				return "RuntimeException";
-			}
-		});
-		handler.setThresholds(Collections.singletonMap("RuntimeException", new Integer(1)));
+		handler.setThresholds(Collections.<Class<? extends Throwable>, Integer> singletonMap(Exception.class, 1));
 		handler.setUseParent(true);
 		// No exception...
 		handler.handleException(context, new RuntimeException("Foo"));

spring-batch-infrastructure/src/test/java/org/springframework/batch/repeat/exception/SimpleLimitExceptionHandlerTests.java

@@ -16,11 +16,17 @@
 
 package org.springframework.batch.repeat.exception;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
-import junit.framework.TestCase;
-
+import org.junit.Before;
+import org.junit.Test;
 import org.springframework.batch.repeat.context.RepeatContextSupport;
 
 /**
@@ -29,77 +35,110 @@ import org.springframework.batch.repeat.context.RepeatContextSupport;
  * @author Robert Kasanicky
  * @author Dave Syer
  */
-public class SimpleLimitExceptionHandlerTests extends TestCase {
+public class SimpleLimitExceptionHandlerTests {
 
 	// object under test
 	private SimpleLimitExceptionHandler handler = new SimpleLimitExceptionHandler();
 
+	@Before
+	public void initializeHandler() throws Exception {
+		handler.afterPropertiesSet();
+	}
+
+	@Test
 	public void testInitializeWithNullContext() throws Throwable {
 		try {
 			handler.handleException(null, new RuntimeException("foo"));
 			fail("Expected IllegalArgumentException");
-		} catch (IllegalArgumentException e) {
+		}
+		catch (IllegalArgumentException e) {
 			// expected
 		}
 	}
 
+	@Test
 	public void testInitializeWithNullContextAndNullException() throws Throwable {
 		try {
 			handler.handleException(null, null);
-		} catch (NullPointerException e) {
+		}
+		catch (IllegalArgumentException e) {
 			// expected;
 		}
 	}
 
-	/**
-	 * Other than nominated exception type should be rethrown, ignoring the exception limit.
-	 * 
-	 * @throws Exception
-	 */
-	public void testNormalExceptionThrown() throws Throwable {
+	@Test
+	public void testDefaultBehaviour() throws Throwable {
 		Throwable throwable = new RuntimeException("foo");
-
-		final int MORE_THAN_ZERO = 1;
-		handler.setLimit(MORE_THAN_ZERO);
-		handler.setExceptionClasses(new Class<?>[] { IllegalArgumentException.class });
-
 		try {
 			handler.handleException(new RepeatContextSupport(null), throwable);
-			fail("Exception swallowed.");
-		} catch (RuntimeException expected) {
+			fail("Exception was swallowed.");
+		}
+		catch (RuntimeException expected) {
 			assertTrue("Exception is rethrown, ignoring the exception limit", true);
 			assertSame(expected, throwable);
 		}
 	}
 
 	/**
-	 * TransactionInvalidException should only be rethrown below the exception limit.
+	 * Other than nominated exception type should be rethrown, ignoring the
+	 * exception limit.
 	 * 
 	 * @throws Exception
 	 */
+	@Test
+	public void testNormalExceptionThrown() throws Throwable {
+		Throwable throwable = new RuntimeException("foo");
+
+		final int MORE_THAN_ZERO = 1;
+		handler.setLimit(MORE_THAN_ZERO);
+		handler.setExceptionClasses(Collections.<Class<? extends Throwable>> singleton(IllegalArgumentException.class));
+		handler.afterPropertiesSet();
+
+		try {
+			handler.handleException(new RepeatContextSupport(null), throwable);
+			fail("Exception was swallowed.");
+		}
+		catch (RuntimeException expected) {
+			assertTrue("Exception is rethrown, ignoring the exception limit", true);
+			assertSame(expected, throwable);
+		}
+	}
+
+	/**
+	 * TransactionInvalidException should only be rethrown below the exception
+	 * limit.
+	 * 
+	 * @throws Exception
+	 */
+	@Test
 	public void testLimitedExceptionTypeNotThrown() throws Throwable {
 		final int MORE_THAN_ZERO = 1;
 		handler.setLimit(MORE_THAN_ZERO);
-		handler.setExceptionClasses(new Class[] {RuntimeException.class} );
+		handler.setExceptionClasses(Collections.<Class<? extends Throwable>> singleton(RuntimeException.class));
+		handler.afterPropertiesSet();
 
 		try {
 			handler.handleException(new RepeatContextSupport(null), new RuntimeException("foo"));
-		} catch (RuntimeException expected) {
+		}
+		catch (RuntimeException expected) {
 			fail("Unexpected exception.");
 		}
 	}
 
 	/**
-	 * TransactionInvalidException should only be rethrown below the exception limit.
+	 * TransactionInvalidException should only be rethrown below the exception
+	 * limit.
 	 * 
 	 * @throws Exception
 	 */
+	@Test
 	public void testLimitedExceptionNotThrownFromSiblings() throws Throwable {
 		Throwable throwable = new RuntimeException("foo");
 
 		final int MORE_THAN_ZERO = 1;
 		handler.setLimit(MORE_THAN_ZERO);
-		handler.setExceptionClasses(new Class[] {RuntimeException.class});
+		handler.setExceptionClasses(Collections.<Class<? extends Throwable>> singleton(RuntimeException.class));
+		handler.afterPropertiesSet();
 
 		RepeatContextSupport parent = new RepeatContextSupport(null);
 
@@ -108,23 +147,27 @@ public class SimpleLimitExceptionHandlerTests extends TestCase {
 			handler.handleException(context, throwable);
 			context = new RepeatContextSupport(parent);
 			handler.handleException(context, throwable);
-		} catch (RuntimeException expected) {
+		}
+		catch (RuntimeException expected) {
 			fail("Unexpected exception.");
 		}
 	}
 
 	/**
-	 * TransactionInvalidException should only be rethrown below the exception limit.
+	 * TransactionInvalidException should only be rethrown below the exception
+	 * limit.
 	 * 
 	 * @throws Exception
 	 */
+	@Test
 	public void testLimitedExceptionThrownFromSiblingsWhenUsingParent() throws Throwable {
 		Throwable throwable = new RuntimeException("foo");
 
 		final int MORE_THAN_ZERO = 1;
 		handler.setLimit(MORE_THAN_ZERO);
-		handler.setExceptionClasses(new Class[] { RuntimeException.class } );
+		handler.setExceptionClasses(Collections.<Class<? extends Throwable>> singleton(RuntimeException.class));
 		handler.setUseParent(true);
+		handler.afterPropertiesSet();
 
 		RepeatContextSupport parent = new RepeatContextSupport(null);
 
@@ -134,19 +177,22 @@ public class SimpleLimitExceptionHandlerTests extends TestCase {
 			context = new RepeatContextSupport(parent);
 			handler.handleException(context, throwable);
 			fail("Expected exception.");
-		} catch (RuntimeException expected) {
+		}
+		catch (RuntimeException expected) {
 			assertSame(throwable, expected);
 		}
 	}
 
 	/**
-	 * TransactionInvalidExceptions are swallowed until the exception limit is exceeded. After the limit is exceeded
-	 * exceptions are rethrown as BatchCriticalExceptions
+	 * Exceptions are swallowed until the exception limit is exceeded. After the
+	 * limit is exceeded exceptions are rethrown
 	 */
+	@Test
 	public void testExceptionNotThrownBelowLimit() throws Throwable {
 
 		final int EXCEPTION_LIMIT = 3;
 		handler.setLimit(EXCEPTION_LIMIT);
+		handler.afterPropertiesSet();
 
 		List<Throwable> throwables = new ArrayList<Throwable>() {
 			{
@@ -165,20 +211,24 @@ public class SimpleLimitExceptionHandlerTests extends TestCase {
 				assertTrue("exceptions up to limit are swallowed", true);
 
 			}
-		} catch (RuntimeException unexpected) {
+		}
+		catch (RuntimeException unexpected) {
 			fail("exception rethrown although exception limit was not exceeded");
 		}
 
 	}
 
 	/**
-	 * TransactionInvalidExceptions are swallowed until the exception limit is exceeded. After the limit is exceeded
-	 * exceptions are rethrown as BatchCriticalExceptions
+	 * TransactionInvalidExceptions are swallowed until the exception limit is
+	 * exceeded. After the limit is exceeded exceptions are rethrown as
+	 * BatchCriticalExceptions
 	 */
+	@Test
 	public void testExceptionThrownAboveLimit() throws Throwable {
 
 		final int EXCEPTION_LIMIT = 3;
 		handler.setLimit(EXCEPTION_LIMIT);
+		handler.afterPropertiesSet();
 
 		List<Throwable> throwables = new ArrayList<Throwable>() {
 			{
@@ -199,7 +249,8 @@ public class SimpleLimitExceptionHandlerTests extends TestCase {
 				assertTrue("exceptions up to limit are swallowed", true);
 
 			}
-		} catch (RuntimeException expected) {
+		}
+		catch (RuntimeException expected) {
 			assertEquals("above exception limit", expected.getMessage());
 		}
 
@@ -208,7 +259,8 @@ public class SimpleLimitExceptionHandlerTests extends TestCase {
 			handler.handleException(context, new RuntimeException("foo"));
 			assertTrue("exceptions up to limit are swallowed", true);
 
-		} catch (RuntimeException expected) {
+		}
+		catch (RuntimeException expected) {
 			assertEquals("foo", expected.getMessage());
 		}
 	}

spring-batch-infrastructure/src/test/java/org/springframework/batch/retry/policy/ExceptionClassifierRetryPolicyTests.java

@@ -18,17 +18,15 @@ package org.springframework.batch.retry.policy;
 
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.Map;
-
 import junit.framework.TestCase;
 
 import org.springframework.batch.retry.RetryContext;
 import org.springframework.batch.retry.RetryPolicy;
-import org.springframework.batch.support.ExceptionClassifierSupport;
+import org.springframework.batch.support.Classifier;
 
 public class ExceptionClassifierRetryPolicyTests extends TestCase {
 
-	ExceptionClassifierRetryPolicy policy = new ExceptionClassifierRetryPolicy();
+	private ExceptionClassifierRetryPolicy policy = new ExceptionClassifierRetryPolicy();
 
 	public void testDefaultPolicies() throws Exception {
 		RetryContext context = policy.open(null);
@@ -36,28 +34,22 @@ public class ExceptionClassifierRetryPolicyTests extends TestCase {
 	}
 
 	public void testTrivialPolicies() throws Exception {
-		policy.setPolicyMap(Collections.singletonMap(ExceptionClassifierSupport.DEFAULT,
-				(RetryPolicy) new MockRetryPolicySupport()));
+		policy.setPolicyMap(Collections.<Class<? extends Throwable>, RetryPolicy> singletonMap(Exception.class,
+				new MockRetryPolicySupport()));
 		RetryContext context = policy.open(null);
 		assertNotNull(context);
 		assertTrue(policy.canRetry(context));
 	}
 
 	public void testNullPolicies() throws Exception {
-		policy.setPolicyMap(new HashMap<String, RetryPolicy>());
-		try {
-			policy.open(null);
-			fail("Expected IllegalArgumentException");
-		}
-		catch (IllegalArgumentException e) {
-			// expected
-		}
+		policy.setPolicyMap(new HashMap<Class<? extends Throwable>, RetryPolicy>());
+		RetryContext context = policy.open(null);
+		assertNotNull(context);
 	}
 
 	public void testNullContext() throws Exception {
-		Map<String, RetryPolicy> map = new HashMap<String, RetryPolicy>();
-		map.put(ExceptionClassifierSupport.DEFAULT, new NeverRetryPolicy());
-		policy.setPolicyMap(map);
+		policy.setPolicyMap(Collections.<Class<? extends Throwable>, RetryPolicy> singletonMap(Exception.class,
+				new NeverRetryPolicy()));
 
 		RetryContext context = policy.open(null);
 		assertNotNull(context);
@@ -67,53 +59,52 @@ public class ExceptionClassifierRetryPolicyTests extends TestCase {
 
 	public void testClassifierOperates() throws Exception {
 
-		Map<String, RetryPolicy> map = new HashMap<String, RetryPolicy>();
-		map.put(ExceptionClassifierSupport.DEFAULT, new AlwaysRetryPolicy());
-		map.put("foo", new NeverRetryPolicy());
-		policy.setPolicyMap(map);
-
 		RetryContext context = policy.open(null);
 		assertNotNull(context);
 
 		assertTrue(policy.canRetry(context));
 		policy.registerThrowable(context, new IllegalArgumentException());
-		assertTrue(policy.canRetry(context));
+		assertFalse(policy.canRetry(context)); // NeverRetryPolicy is the
+		// default
 
-		policy.setExceptionClassifier(new ExceptionClassifierSupport() {
-			public String classify(Throwable throwable) {
+		policy.setExceptionClassifier(new Classifier<Throwable, RetryPolicy>() {
+			public RetryPolicy classify(Throwable throwable) {
 				if (throwable != null) {
-					return "foo";
+					return new AlwaysRetryPolicy();
 				}
-				return super.classify(throwable);
+				return new NeverRetryPolicy();
 			}
 		});
 
 		// The context saves the classifier, so changing it now has no effect
-		assertTrue(policy.canRetry(context));
+		assertFalse(policy.canRetry(context));
 		policy.registerThrowable(context, new IllegalArgumentException());
-		assertTrue(policy.canRetry(context));
+		assertFalse(policy.canRetry(context));
 
 		// But now the classifier will be active in the new context...
 		context = policy.open(null);
 		assertTrue(policy.canRetry(context));
 		policy.registerThrowable(context, new IllegalArgumentException());
-		assertFalse(policy.canRetry(context));
+		assertTrue(policy.canRetry(context));
 
 	}
 
 	int count = 0;
 
 	public void testClose() throws Exception {
-		policy.setPolicyMap(Collections.singletonMap(ExceptionClassifierSupport.DEFAULT,
-				(RetryPolicy) new MockRetryPolicySupport() {
+		policy.setExceptionClassifier(new Classifier<Throwable, RetryPolicy>() {
+			public RetryPolicy classify(Throwable throwable) {
+				return new MockRetryPolicySupport() {
 					public void close(RetryContext context) {
 						count++;
 					}
-				}));
+				};
+			}
+		});
 		RetryContext context = policy.open(null);
 
 		// The mapped (child) policy hasn't been used yet, so if we close now
-		// we don't incur the possible expense of ceating the child context.
+		// we don't incur the possible expense of creating the child context.
 		policy.close(context);
 		assertEquals(0, count); // not classified yet
 		// This forces a child context to be created and the child policy is

spring-batch-infrastructure/src/test/java/org/springframework/batch/retry/policy/FatalExceptionRetryPolicyTests.java

@@ -16,7 +16,8 @@
 
 package org.springframework.batch.retry.policy;
 
-import java.util.HashSet;
+import java.util.Arrays;
+import java.util.List;
 
 import junit.framework.TestCase;
 
@@ -38,12 +39,10 @@ public class FatalExceptionRetryPolicyTests extends TestCase {
 		retryTemplate.setRetryPolicy(policy);
 
 		// ...but make sure certain exceptions are fatal
-		policy.setFatalExceptionClasses(new HashSet<Class<? extends Throwable>>() {
-			{
-				add(IllegalArgumentException.class);
-				add(IllegalStateException.class);
-			}
-		});
+		@SuppressWarnings("unchecked")
+		List<Class<? extends Throwable>> list = Arrays.<Class<? extends Throwable>> asList(IllegalArgumentException.class,
+				IllegalStateException.class);
+		policy.setFatalExceptionClasses(list);
 		RecoveryCallback<String> recoveryCallback = new RecoveryCallback<String>() {
 			public String recover(RetryContext context) throws Exception {
 				return "bar";
@@ -71,13 +70,11 @@ public class FatalExceptionRetryPolicyTests extends TestCase {
 		SimpleRetryPolicy policy = new SimpleRetryPolicy(3);
 		retryTemplate.setRetryPolicy(policy);
 
-		policy.setFatalExceptionClasses(new HashSet<Class<? extends Throwable>>() {
-			{
-				add(IllegalArgumentException.class);
-				add(IllegalStateException.class);
-			}
-		});
-		RecoveryCallback<String>recoveryCallback = new RecoveryCallback<String>() {
+		@SuppressWarnings("unchecked")
+		List<Class<? extends Throwable>> list = Arrays.<Class<? extends Throwable>> asList(
+				IllegalArgumentException.class, IllegalStateException.class);
+		policy.setFatalExceptionClasses(list);
+		RecoveryCallback<String> recoveryCallback = new RecoveryCallback<String>() {
 			public String recover(RetryContext context) throws Exception {
 				return "bar";
 			}

spring-batch-infrastructure/src/test/java/org/springframework/batch/retry/support/RetryTemplateTests.java

@@ -22,19 +22,21 @@ import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.fail;
 
-import java.util.HashSet;
+import java.util.Collections;
 
 import org.junit.Test;
 import org.springframework.batch.retry.ExhaustedRetryException;
 import org.springframework.batch.retry.RecoveryCallback;
 import org.springframework.batch.retry.RetryCallback;
 import org.springframework.batch.retry.RetryContext;
+import org.springframework.batch.retry.RetryState;
 import org.springframework.batch.retry.backoff.BackOffContext;
 import org.springframework.batch.retry.backoff.BackOffInterruptedException;
 import org.springframework.batch.retry.backoff.BackOffPolicy;
 import org.springframework.batch.retry.backoff.StatelessBackOffPolicy;
 import org.springframework.batch.retry.policy.NeverRetryPolicy;
 import org.springframework.batch.retry.policy.SimpleRetryPolicy;
+import org.springframework.batch.support.BinaryExceptionClassifier;
 
 /**
  * @author Rob Harrop
@@ -71,7 +73,7 @@ public class RetryTemplateTests {
 			}
 		});
 		assertEquals(2, callback.attempts);
-		assertEquals(value, result);		
+		assertEquals(value, result);
 	}
 
 	@Test
@@ -117,16 +119,28 @@ public class RetryTemplateTests {
 		assertEquals(attempts, callback.attempts);
 	}
 
+	@Test
+	public void testRollbackClassifierOverridesRetryPolicy() throws Exception {
+		MockRetryCallback callback = new MockRetryCallback();
+		int attempts = 3;
+		callback.setAttemptsBeforeSuccess(attempts);
+		callback.setExceptionToThrow(new IllegalArgumentException());
+
+		RetryTemplate retryTemplate = new RetryTemplate();
+		retryTemplate.setRetryPolicy(new SimpleRetryPolicy(attempts));
+		BinaryExceptionClassifier classifier = new BinaryExceptionClassifier(Collections
+				.<Class<? extends Throwable>> singleton(IllegalArgumentException.class), false);
+		retryTemplate.setRollbackClassifier(classifier);
+		retryTemplate.execute(callback, new RetryState("foo"));
+		assertEquals(attempts, callback.attempts);
+	}
+
 	@Test
 	public void testSetExceptions() throws Exception {
 		RetryTemplate template = new RetryTemplate();
 		SimpleRetryPolicy policy = new SimpleRetryPolicy();
 		template.setRetryPolicy(policy);
-		policy.setRetryableExceptionClasses(new HashSet<Class<? extends Throwable>>() {
-			{
-				add(RuntimeException.class);
-			}
-		});
+		policy.setRetryableExceptionClasses(Collections.<Class<? extends Throwable>> singleton(RuntimeException.class));
 
 		int attempts = 3;
 
@@ -243,7 +257,7 @@ public class RetryTemplateTests {
 			assertEquals("foo", e.getMessage());
 		}
 	}
-	
+
 	private static class MockRetryCallback implements RetryCallback<Object> {
 
 		private int attempts;

spring-batch-infrastructure/src/test/java/org/springframework/batch/retry/support/StatefulRecoveryRetryTests.java

@@ -23,6 +23,7 @@ import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 import org.junit.Test;
@@ -39,6 +40,8 @@ import org.springframework.batch.retry.RetryState;
 import org.springframework.batch.retry.policy.MapRetryContextCache;
 import org.springframework.batch.retry.policy.NeverRetryPolicy;
 import org.springframework.batch.retry.policy.SimpleRetryPolicy;
+import org.springframework.batch.support.BinaryExceptionClassifier;
+import org.springframework.dao.DataAccessException;
 
 public class StatefulRecoveryRetryTests {
 
@@ -121,6 +124,37 @@ public class StatefulRecoveryRetryTests {
 		assertEquals(input, list.get(0));
 	}
 
+	@Test
+	public void testSwitchToStatelessForNoRollback() throws Exception {
+		retryTemplate.setRetryPolicy(new SimpleRetryPolicy(1));
+		// Roll back for these:
+		BinaryExceptionClassifier classifier = new BinaryExceptionClassifier(Collections
+				.<Class<? extends Throwable>> singleton(DataAccessException.class));
+		// ...but not these:
+		assertFalse(classifier.classify(new RuntimeException()));
+		retryTemplate.setRollbackClassifier(classifier);
+		final String input = "foo";
+		RetryState state = new RetryState(input);
+		RetryCallback<String> callback = new RetryCallback<String>() {
+			public String doWithRetry(RetryContext context) throws Exception {
+				throw new RuntimeException("Barf!");
+			}
+		};
+		RecoveryCallback<String> recoveryCallback = new RecoveryCallback<String>() {
+			public String recover(RetryContext context) {
+				count++;
+				list.add(input);
+				return input;
+			}
+		};
+		Object result = null;
+		// On the second retry, the recovery path is taken...
+		result = retryTemplate.execute(callback, recoveryCallback, state);
+		assertEquals(input, result); // default result is the item
+		assertEquals(1, count);
+		assertEquals(input, list.get(0));
+	}
+
 	@Test
 	public void testExhaustedClearsHistoryAfterLastAttempt() throws Exception {
 		RetryPolicy retryPolicy = new SimpleRetryPolicy(1);

spring-batch-infrastructure/src/test/java/org/springframework/batch/support/BinaryExceptionClassifierTests.java

@@ -16,8 +16,7 @@
 
 package org.springframework.batch.support;
 
-import java.util.HashSet;
-
+import java.util.Collections;
 import junit.framework.TestCase;
 
 public class BinaryExceptionClassifierTests extends TestCase {
@@ -25,19 +24,36 @@ public class BinaryExceptionClassifierTests extends TestCase {
 	BinaryExceptionClassifier classifier = new BinaryExceptionClassifier();
 
 	public void testClassifyNullIsDefault() {
-		assertTrue(classifier.isDefault(null));
+		assertFalse(classifier.classify(null));
+	}
+
+	public void testFalseIsDefault() {
+		assertFalse(classifier.getDefault());
+	}
+
+	public void testDefaultProvided() {
+		classifier = new BinaryExceptionClassifier(true);
+		assertTrue(classifier.getDefault());
 	}
 
 	public void testClassifyRandomException() {
-		assertTrue(classifier.isDefault(new IllegalStateException("foo")));
+		assertFalse(classifier.classify(new IllegalStateException("foo")));
 	}
 
 	public void testClassifyExactMatch() {
-		classifier.setExceptionClasses(new HashSet<Class<? extends Throwable>>() {
-			{
-				add(IllegalStateException.class);
-			}
-		});
-		assertEquals(false, classifier.isDefault(new IllegalStateException("Foo")));
+		classifier.setTypes(Collections.<Class<? extends Throwable>> singleton(IllegalStateException.class));
+		assertTrue(classifier.classify(new IllegalStateException("Foo")));
+	}
+
+	public void testTypesProvidedInConstructor() {
+		classifier = new BinaryExceptionClassifier(Collections
+				.<Class<? extends Throwable>> singleton(IllegalStateException.class));
+		assertTrue(classifier.classify(new IllegalStateException("Foo")));
+	}
+
+	public void testTypesProvidedInConstructorWithNonDefault() {
+		classifier = new BinaryExceptionClassifier(Collections
+				.<Class<? extends Throwable>> singleton(IllegalStateException.class), false);
+		assertFalse(classifier.classify(new IllegalStateException("Foo")));
 	}
 }

spring-batch-infrastructure/src/test/java/org/springframework/batch/support/ClassifierSupportTests.java

@@ -16,20 +16,18 @@
 
 package org.springframework.batch.support;
 
-import org.springframework.batch.support.ExceptionClassifierSupport;
-
 import junit.framework.TestCase;
 
-public class ExceptionClassifierSupportTests extends TestCase {
+public class ClassifierSupportTests extends TestCase {
 
 	public void testClassifyNullIsDefault() {
-		ExceptionClassifierSupport classifier = new ExceptionClassifierSupport();
-		assertEquals(classifier.classify(null), classifier.getDefault());
+		ClassifierSupport<String,String> classifier = new ClassifierSupport<String,String>("foo");
+		assertEquals(classifier.classify(null), "foo");
 	}
 
 	public void testClassifyRandomException() {
-		ExceptionClassifierSupport classifier = new ExceptionClassifierSupport();
-		assertEquals(classifier.classify(new IllegalStateException("Foo")), classifier.getDefault());
+		ClassifierSupport<Throwable,String> classifier = new ClassifierSupport<Throwable,String>("foo");
+		assertEquals(classifier.classify(new IllegalStateException("Foo")), classifier.classify(null));
 	}
 
 }

spring-batch-infrastructure/src/test/java/org/springframework/batch/support/SubclassExceptionClassifierTests.java

@@ -16,49 +16,69 @@
 
 package org.springframework.batch.support;
 
-import java.util.LinkedHashMap;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
 
-import junit.framework.TestCase;
+import java.util.Collections;
+import java.util.HashMap;
 
-public class SubclassExceptionClassifierTests extends TestCase {
+import org.junit.Test;
 
-	SubclassExceptionClassifier classifier = new SubclassExceptionClassifier();
+public class SubclassExceptionClassifierTests {
 
+	SubclassClassifier<Throwable, String> classifier = new SubclassClassifier<Throwable, String>();
+
+	@Test
 	public void testClassifyNullIsDefault() {
 		assertEquals(classifier.classify(null), classifier.getDefault());
 	}
 
+	@Test
+	public void testClassifyNull() {
+		assertNull(classifier.classify(null));
+	}
+
+	@Test
+	public void testClassifyNullNonDefault() {
+		classifier = new SubclassClassifier<Throwable, String>("foo");
+		assertEquals("foo", classifier.classify(null));
+	}
+
+	@Test
 	public void testClassifyRandomException() {
-		assertEquals(classifier.classify(new IllegalStateException("Foo")), classifier.getDefault());
+		assertNull(classifier.classify(new IllegalStateException("Foo")));
 	}
 
+	@Test
 	public void testClassifyExactMatch() {
-		classifier.setTypeMap(new LinkedHashMap<Class<? extends Throwable>, String>() {{
-			put(IllegalStateException.class, "foo");
-		}});
+		classifier.setTypeMap(Collections.<Class<? extends Throwable>, String> singletonMap(
+				IllegalStateException.class, "foo"));
 		assertEquals("foo", classifier.classify(new IllegalStateException("Foo")));
 	}
 
+	@Test
 	public void testClassifySubclassMatch() {
-		classifier.setTypeMap(new LinkedHashMap<Class<? extends Throwable>, String>() {{
-			put(RuntimeException.class, "foo");
-		}});
+		classifier.setTypeMap(Collections.<Class<? extends Throwable>, String> singletonMap(RuntimeException.class,
+				"foo"));
 		assertEquals("foo", classifier.classify(new IllegalStateException("Foo")));
 	}
 
+	@Test
 	public void testClassifySuperclassDoesNotMatch() {
-		classifier.setTypeMap(new LinkedHashMap<Class<? extends Throwable>, String>() {{
-			put(IllegalStateException.class, "foo");
-		}});
+		classifier.setTypeMap(Collections.<Class<? extends Throwable>, String> singletonMap(
+				IllegalStateException.class, "foo"));
 		assertEquals(classifier.getDefault(), classifier.classify(new RuntimeException("Foo")));
 	}
 
+	@Test
 	public void testClassifyAncestorMatch() {
-		classifier.setTypeMap(new LinkedHashMap<Class<? extends Throwable>, String>() {{
-			put(Exception.class, "bar");
-			put(IllegalArgumentException.class, "foo");
-			put(RuntimeException.class, "bucket");
-		}});
-		assertEquals("bucket", classifier.classify(new IllegalStateException("Foo")));
+		classifier.setTypeMap(new HashMap<Class<? extends Throwable>, String>() {
+			{
+				put(Exception.class, "foo");
+				put(IllegalArgumentException.class, "bar");
+				put(RuntimeException.class, "spam");
+			}
+		});
+		assertEquals("spam", classifier.classify(new IllegalStateException("Foo")));
 	}
 }