diff --git a/infrastructure/src/test/java/org/springframework/batch/io/file/support/DefaultFlatFileInputSourceTests.java b/infrastructure/src/test/java/org/springframework/batch/io/file/support/DefaultFlatFileInputSourceTests.java
index cd8789c7a..d74acd31c 100644
--- a/infrastructure/src/test/java/org/springframework/batch/io/file/support/DefaultFlatFileInputSourceTests.java
+++ b/infrastructure/src/test/java/org/springframework/batch/io/file/support/DefaultFlatFileInputSourceTests.java
@@ -158,6 +158,7 @@ public class DefaultFlatFileInputSourceTests extends TestCase {
 	public void testRestartWithNullReader() throws Exception {
 		inputSource = new DefaultFlatFileInputSource();
 		inputSource.setResource(getInputResource(TEST_STRING));
+		inputSource.setFieldSetMapper(fieldSetMapper);
 		// do not open the template...
 		inputSource.restoreFrom(inputSource.getRestartData());
 		assertEquals("[FlatFileInputTemplate-TestData]", inputSource.read().toString());
diff --git a/infrastructure/src/test/java/org/springframework/batch/io/file/support/SimpleFlatFileInputSourceTests.java b/infrastructure/src/test/java/org/springframework/batch/io/file/support/SimpleFlatFileInputSourceTests.java
index 191f6a51f..6fc2fe5f8 100644
--- a/infrastructure/src/test/java/org/springframework/batch/io/file/support/SimpleFlatFileInputSourceTests.java
+++ b/infrastructure/src/test/java/org/springframework/batch/io/file/support/SimpleFlatFileInputSourceTests.java
@@ -139,12 +139,14 @@ public class SimpleFlatFileInputSourceTests extends TestCase {
 	public void testReadBeforeOpen() throws Exception {
 		inputSource = new SimpleFlatFileInputSource();
 		inputSource.setResource(getInputResource(TEST_STRING));
+		inputSource.setFieldSetMapper(fieldSetMapper);
 		assertEquals("[FlatFileInputTemplate-TestData]", inputSource.read().toString());
 	}
 
 	public void testCloseBeforeOpen() throws Exception {
 		inputSource = new SimpleFlatFileInputSource();
 		inputSource.setResource(getInputResource(TEST_STRING));
+		inputSource.setFieldSetMapper(fieldSetMapper);
 		inputSource.close();
 		// The open still happens automatically on a read...
 		assertEquals("[FlatFileInputTemplate-TestData]", inputSource.read().toString());
@@ -181,6 +183,7 @@ public class SimpleFlatFileInputSourceTests extends TestCase {
 		inputSource = new SimpleFlatFileInputSource();
 		inputSource.setEncoding("UTF-8");
 		inputSource.setResource(getInputResource(TEST_STRING));
+		inputSource.setFieldSetMapper(fieldSetMapper);
 		testRead();
 	}