Fixed StepScope not being handled correctly - creating too many auto-proxies for the same stepExecution.

spring-batch-core/src/test/java/org/springframework/batch/core/partition/RestartIntegrationTests.java

@@ -87,7 +87,6 @@ public class RestartIntegrationTests {
 		// Two attempts
 		assertEquals(2, afterMaster-beforeMaster);
 		// One failure and two successes
-		System.out.println(afterPartition);
 		assertEquals(3, afterPartition-beforePartition);
 
 	}

spring-batch-core/src/test/java/org/springframework/batch/core/scope/StepContextRepeatCallbackTests.java

@@ -19,6 +19,7 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import org.junit.After;
 import org.junit.Test;
 import org.springframework.batch.core.ExitStatus;
 import org.springframework.batch.core.JobExecution;
@@ -35,6 +36,11 @@ public class StepContextRepeatCallbackTests {
 	private StepExecution stepExecution = new StepExecution("foo", new JobExecution(0L), 123L);
 	private boolean addedAttribute = false;
 	private boolean removedAttribute = false;
+	
+	@After
+	public void cleanUpStepContext() {
+		StepSynchronizationManager.close();
+	}
 
 	@Test
 	public void testDoInIteration() throws Exception {
@@ -51,6 +57,7 @@ public class StepContextRepeatCallbackTests {
 
 	@Test
 	public void testUnfinishedWork() throws Exception {
+		StepSynchronizationManager.register(stepExecution);
 		StepContextRepeatCallback callback = new StepContextRepeatCallback(stepExecution) {
 			@Override
 			public RepeatStatus doInStepContext(RepeatContext context, StepContext stepContext) throws Exception {

spring-batch-core/src/test/java/org/springframework/batch/core/scope/StepScopeIntegrationTests.java

@@ -38,6 +38,7 @@ public class StepScopeIntegrationTests {
 	@Before
 	@After
 	public void start() {
+		StepSynchronizationManager.close();
 		TestStep.reset();
 	}
 		

spring-batch-core/src/test/java/org/springframework/batch/core/scope/StepScopeTests.java

@@ -18,6 +18,7 @@ package org.springframework.batch.core.scope;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -41,11 +42,13 @@ public class StepScopeTests {
 
 	private StepScope scope = new StepScope();
 
-	private StepContext context = new StepContext(new StepExecution("foo", new JobExecution(0L)));
+	private StepExecution stepExecution = new StepExecution("foo", new JobExecution(0L));
+
+	private StepContext context;
 
 	@Before
 	public void setUp() throws Exception {
-		StepSynchronizationManager.register(context);
+		context = StepSynchronizationManager.register(stepExecution);
 	}
 
 	@After
@@ -98,8 +101,7 @@ public class StepScopeTests {
 	@Test
 	public void testGetWithSomethingAlreadyInParentContext() {
 		context.setAttribute("foo", "bar");
-		StepContext context = new StepContext(new StepExecution("bar", new JobExecution(0L)));
-		StepSynchronizationManager.register(context);
+		StepContext context = StepSynchronizationManager.register(new StepExecution("bar", new JobExecution(0L)));
 		Object value = scope.get("foo", new ObjectFactory() {
 			public Object getObject() throws BeansException {
 				return "spam";
@@ -111,6 +113,13 @@ public class StepScopeTests {
 		assertEquals("bar", scope.get("foo", null));
 	}
 
+	@Test
+	public void testParentContextWithSameStepExecution() {
+		context.setAttribute("foo", "bar");
+		StepContext other = StepSynchronizationManager.register(stepExecution);
+		assertSame(other, context);
+	}
+
 	@Test
 	public void testGetConversationId() {
 		String id = scope.getConversationId();

spring-batch-core/src/test/java/org/springframework/batch/core/scope/StepSynchronizationManagerTests.java

@@ -28,20 +28,19 @@ public class StepSynchronizationManagerTests {
 	@Test
 	public void testGetContext() {
 		assertNull(StepSynchronizationManager.getContext());
-		StepSynchronizationManager.register(new StepContext(stepExecution));
+		StepSynchronizationManager.register(stepExecution);
 		assertNotNull(StepSynchronizationManager.getContext());
 	}
 
 	@Test
 	public void testClose() {
-		StepContext context = new StepContext(stepExecution);
 		final List<String> list = new ArrayList<String>();
+		StepContext context = StepSynchronizationManager.register(stepExecution);
 		context.registerDestructionCallback("foo", new Runnable() {
 			public void run() {
 				list.add("foo");
 			}
 		});
-		StepSynchronizationManager.register(context);
 		StepSynchronizationManager.close();
 		assertNull(StepSynchronizationManager.getContext());
 		assertEquals(0, list.size());
@@ -49,14 +48,13 @@ public class StepSynchronizationManagerTests {
 
 	@Test
 	public void testRelease() {
-		StepContext context = new StepContext(stepExecution);
+		StepContext context = StepSynchronizationManager.register(stepExecution);
 		final List<String> list = new ArrayList<String>();
 		context.registerDestructionCallback("foo", new Runnable() {
 			public void run() {
 				list.add("foo");
 			}
 		});
-		StepSynchronizationManager.register(context);
 		// On release we expect the destruction callbacks to be called
 		StepSynchronizationManager.release();
 		assertNull(StepSynchronizationManager.getContext());
@@ -72,9 +70,8 @@ public class StepSynchronizationManagerTests {
 
 	@Test
 	public void testRegisterTwice() {
-		StepContext context = new StepContext(stepExecution);
-		StepSynchronizationManager.register(context);
-		StepSynchronizationManager.register(context);
+		StepSynchronizationManager.register(stepExecution);
+		StepSynchronizationManager.register(stepExecution);
 		StepSynchronizationManager.close();
 		// if someone registers you have to assume they are going to close, so
 		// the last thing you want is for the close to remove another context

spring-batch-core/src/test/java/org/springframework/batch/core/scope/TestStep.java

@@ -7,13 +7,13 @@ import org.springframework.batch.core.StepExecution;
 public class TestStep implements Step {
 
 	private static StepContext context;
-	
+
 	private Collaborator collaborator;
-	
+
 	public void setCollaborator(Collaborator collaborator) {
 		this.collaborator = collaborator;
 	}
-	
+
 	public static StepContext getContext() {
 		return context;
 	}
@@ -24,7 +24,9 @@ public class TestStep implements Step {
 
 	public void execute(StepExecution stepExecution) throws JobInterruptedException {
 		context = StepSynchronizationManager.getContext();
-		context.setAttribute("collaborator", collaborator.getName());
+		if (context != null) {
+			context.setAttribute("collaborator", collaborator.getName());
+		}
 	}
 
 	public String getName() {