BATCH-1309: Added a 'strict' property to ExecutionContextPromotionListener to ensure the key is present

spring-batch-core/src/main/java/org/springframework/batch/core/listener/ExecutionContextPromotionListener.java

@@ -45,6 +45,8 @@ public class ExecutionContextPromotionListener extends StepExecutionListenerSupp
 
 	private String[] statuses = new String[] { ExitStatus.COMPLETED.getExitCode() };
 
+	private boolean strict = false;
+
 	public ExitStatus afterStep(StepExecution stepExecution) {
 		ExecutionContext stepContext = stepExecution.getExecutionContext();
 		ExecutionContext jobContext = stepExecution.getJobExecution().getExecutionContext();
@@ -52,6 +54,10 @@ public class ExecutionContextPromotionListener extends StepExecutionListenerSupp
 		for (String statusPattern : statuses) {
 			if (PatternMatcher.match(statusPattern, exitCode)) {
 				for (String key : keys) {
+					if (strict) {
+						Assert.isTrue(stepContext.containsKey(key), "The key [" + key
+								+ "] was not found in the Step's ExecutionContext.");
+					}
 					jobContext.put(key, stepContext.get(key));
 				}
 				break;
@@ -84,4 +90,15 @@ public class ExecutionContextPromotionListener extends StepExecutionListenerSupp
 	public void setStatuses(String[] statuses) {
 		this.statuses = statuses;
 	}
+
+	/**
+	 * If set to TRUE, the listener will throw an exception if any 'key' is not
+	 * found in the Step {@link ExecutionContext}. FALSE by default.
+	 * 
+	 * @param strict
+	 */
+	public void setStrict(boolean strict) {
+		this.strict = strict;
+	}
+
 }

spring-batch-core/src/test/java/org/springframework/batch/core/listener/ExecutionContextPromotionListenerTests.java

@@ -174,6 +174,35 @@ public class ExecutionContextPromotionListenerTests {
 		assertFalse(jobExecution.getExecutionContext().containsKey(key2));
 	}
 
+	/**
+	 * CONDITION: strict = true. keys = {key, key2}. Only {key} exists in the
+	 * ExecutionContext.
+	 * 
+	 * EXPECTED: IllegalArgumentException
+	 */
+	@Test(expected = IllegalArgumentException.class)
+	public void promoteEntries_keyNotFound_strict() throws Exception {
+		ExecutionContextPromotionListener listener = new ExecutionContextPromotionListener();
+		listener.setStrict(true);
+
+		JobExecution jobExecution = new JobExecution(1L);
+		StepExecution stepExecution = jobExecution.createStepExecution("step1");
+		stepExecution.setExitStatus(ExitStatus.COMPLETED);
+
+		Assert.state(jobExecution.getExecutionContext().isEmpty());
+		Assert.state(stepExecution.getExecutionContext().isEmpty());
+
+		stepExecution.getExecutionContext().putString(key, value);
+
+		listener.setKeys(new String[] { key, key2 });
+		listener.afterPropertiesSet();
+
+		listener.afterStep(stepExecution);
+
+		assertEquals(value, jobExecution.getExecutionContext().getString(key));
+		assertFalse(jobExecution.getExecutionContext().containsKey(key2));
+	}
+
 	/**
 	 * CONDITION: keys = NULL
 	 *