From 5bb8292e254641e3cc61d5c82ad8be8b66b53b44 Mon Sep 17 00:00:00 2001 From: Daniel Kulp Date: Tue, 22 Nov 2016 13:31:19 -0500 Subject: [PATCH] [BEAM-1034] Clean up tmp area in tests --- .../sorter/BufferedExternalSorter.java | 6 +- .../sdk/extensions/sorter/ExternalSorter.java | 6 +- .../sorter/BufferedExternalSorterTest.java | 58 ++++++++++++++++--- .../extensions/sorter/ExternalSorterTest.java | 53 ++++++++++++++--- 4 files changed, 103 insertions(+), 20 deletions(-) diff --git a/sdks/java/extensions/sorter/src/main/java/org/apache/beam/sdk/extensions/sorter/BufferedExternalSorter.java b/sdks/java/extensions/sorter/src/main/java/org/apache/beam/sdk/extensions/sorter/BufferedExternalSorter.java index 0f89e301467a..1dfd3395365a 100644 --- a/sdks/java/extensions/sorter/src/main/java/org/apache/beam/sdk/extensions/sorter/BufferedExternalSorter.java +++ b/sdks/java/extensions/sorter/src/main/java/org/apache/beam/sdk/extensions/sorter/BufferedExternalSorter.java @@ -35,12 +35,13 @@ public static class Options implements Serializable { private int memoryMB = 100; /** Sets the path to a temporary location where the sorter writes intermediate files. */ - public void setTempLocation(String tempLocation) { + public Options setTempLocation(String tempLocation) { checkArgument( !tempLocation.startsWith("gs://"), "BufferedExternalSorter does not support GCS temporary location"); this.tempLocation = tempLocation; + return this; } /** Returns the configured temporary location. */ @@ -52,9 +53,10 @@ public String getTempLocation() { * Sets the size of the memory buffer in megabytes. This controls both the buffer for initial in * memory sorting and the buffer used when external sorting. Must be greater than zero. */ - public void setMemoryMB(int memoryMB) { + public Options setMemoryMB(int memoryMB) { checkArgument(memoryMB > 0, "memoryMB must be greater than zero"); this.memoryMB = memoryMB; + return this; } /** Returns the configured size of the memory buffer. */ diff --git a/sdks/java/extensions/sorter/src/main/java/org/apache/beam/sdk/extensions/sorter/ExternalSorter.java b/sdks/java/extensions/sorter/src/main/java/org/apache/beam/sdk/extensions/sorter/ExternalSorter.java index 3cf0cc00e73a..beef1eebc4b5 100644 --- a/sdks/java/extensions/sorter/src/main/java/org/apache/beam/sdk/extensions/sorter/ExternalSorter.java +++ b/sdks/java/extensions/sorter/src/main/java/org/apache/beam/sdk/extensions/sorter/ExternalSorter.java @@ -67,12 +67,13 @@ public static class Options implements Serializable { private int memoryMB = 100; /** Sets the path to a temporary location where the sorter writes intermediate files. */ - public void setTempLocation(String tempLocation) { + public Options setTempLocation(String tempLocation) { if (tempLocation.startsWith("gs://")) { throw new IllegalArgumentException("Sorter doesn't support GCS temporary location."); } this.tempLocation = tempLocation; + return this; } /** Returns the configured temporary location. */ @@ -81,9 +82,10 @@ public String getTempLocation() { } /** Sets the size of the memory buffer in megabytes. */ - public void setMemoryMB(int memoryMB) { + public Options setMemoryMB(int memoryMB) { checkArgument(memoryMB > 0, "memoryMB must be greater than zero"); this.memoryMB = memoryMB; + return this; } /** Returns the configured size of the memory buffer. */ diff --git a/sdks/java/extensions/sorter/src/test/java/org/apache/beam/sdk/extensions/sorter/BufferedExternalSorterTest.java b/sdks/java/extensions/sorter/src/test/java/org/apache/beam/sdk/extensions/sorter/BufferedExternalSorterTest.java index 63dbedfaafb7..8c108ebf8024 100644 --- a/sdks/java/extensions/sorter/src/test/java/org/apache/beam/sdk/extensions/sorter/BufferedExternalSorterTest.java +++ b/sdks/java/extensions/sorter/src/test/java/org/apache/beam/sdk/extensions/sorter/BufferedExternalSorterTest.java @@ -27,9 +27,17 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import java.io.IOException; +import java.nio.file.FileVisitResult; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.SimpleFileVisitor; +import java.nio.file.attribute.BasicFileAttributes; import java.util.Arrays; import org.apache.beam.sdk.extensions.sorter.SorterTestUtils.SorterGenerator; import org.apache.beam.sdk.values.KV; +import org.junit.AfterClass; +import org.junit.BeforeClass; import org.junit.Rule; import org.junit.Test; import org.junit.rules.ExpectedException; @@ -40,6 +48,29 @@ @RunWith(JUnit4.class) public class BufferedExternalSorterTest { @Rule public ExpectedException thrown = ExpectedException.none(); + static Path tmpLocation; + + @BeforeClass + public static void setupTempDir() throws IOException { + tmpLocation = Files.createTempDirectory("tmp"); + } + + @AfterClass + public static void cleanupTempDir() throws IOException { + Files.walkFileTree(tmpLocation, new SimpleFileVisitor() { + @Override + public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException { + Files.delete(file); + return FileVisitResult.CONTINUE; + } + + @Override + public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOException { + Files.delete(dir); + return FileVisitResult.CONTINUE; + } + }); + } @SuppressWarnings("unchecked") @Test @@ -106,25 +137,29 @@ public void testFallback() throws Exception { @Test public void testEmpty() throws Exception { - SorterTestUtils.testEmpty(BufferedExternalSorter.create(new BufferedExternalSorter.Options())); + SorterTestUtils.testEmpty(BufferedExternalSorter.create(new BufferedExternalSorter.Options() + .setTempLocation(tmpLocation.toString()))); } @Test public void testSingleElement() throws Exception { SorterTestUtils.testSingleElement( - BufferedExternalSorter.create(new BufferedExternalSorter.Options())); + BufferedExternalSorter.create(new BufferedExternalSorter.Options() + .setTempLocation(tmpLocation.toString()))); } @Test public void testEmptyKeyValueElement() throws Exception { SorterTestUtils.testEmptyKeyValueElement( - BufferedExternalSorter.create(new BufferedExternalSorter.Options())); + BufferedExternalSorter.create(new BufferedExternalSorter.Options() + .setTempLocation(tmpLocation.toString()))); } @Test public void testMultipleIterations() throws Exception { SorterTestUtils.testMultipleIterations( - BufferedExternalSorter.create(new BufferedExternalSorter.Options())); + BufferedExternalSorter.create(new BufferedExternalSorter.Options() + .setTempLocation(tmpLocation.toString()))); } @Test @@ -133,7 +168,8 @@ public void testManySortersFewRecords() throws Exception { new SorterGenerator() { @Override public Sorter generateSorter() throws Exception { - return BufferedExternalSorter.create(new BufferedExternalSorter.Options()); + return BufferedExternalSorter.create(new BufferedExternalSorter.Options() + .setTempLocation(tmpLocation.toString())); } }, 1000000, @@ -146,7 +182,8 @@ public void testOneSorterManyRecords() throws Exception { new SorterGenerator() { @Override public Sorter generateSorter() throws Exception { - return BufferedExternalSorter.create(new BufferedExternalSorter.Options()); + return BufferedExternalSorter.create(new BufferedExternalSorter.Options() + .setTempLocation(tmpLocation.toString())); } }, 1, @@ -156,14 +193,16 @@ public Sorter generateSorter() throws Exception { @Test public void testAddAfterSort() throws Exception { SorterTestUtils.testAddAfterSort( - BufferedExternalSorter.create(new BufferedExternalSorter.Options()), thrown); + BufferedExternalSorter.create(new BufferedExternalSorter.Options() + .setTempLocation(tmpLocation.toString())), thrown); fail(); } @Test public void testSortTwice() throws Exception { SorterTestUtils.testSortTwice( - BufferedExternalSorter.create(new BufferedExternalSorter.Options()), thrown); + BufferedExternalSorter.create(new BufferedExternalSorter.Options() + .setTempLocation(tmpLocation.toString())), thrown); fail(); } @@ -171,7 +210,8 @@ public void testSortTwice() throws Exception { public void testNegativeMemory() throws Exception { thrown.expect(IllegalArgumentException.class); thrown.expectMessage("memoryMB must be greater than zero"); - BufferedExternalSorter.Options options = new BufferedExternalSorter.Options(); + BufferedExternalSorter.Options options = new BufferedExternalSorter.Options() + .setTempLocation(tmpLocation.toString()); options.setMemoryMB(-1); } } diff --git a/sdks/java/extensions/sorter/src/test/java/org/apache/beam/sdk/extensions/sorter/ExternalSorterTest.java b/sdks/java/extensions/sorter/src/test/java/org/apache/beam/sdk/extensions/sorter/ExternalSorterTest.java index 9232b6229957..bcfbdade94a7 100644 --- a/sdks/java/extensions/sorter/src/test/java/org/apache/beam/sdk/extensions/sorter/ExternalSorterTest.java +++ b/sdks/java/extensions/sorter/src/test/java/org/apache/beam/sdk/extensions/sorter/ExternalSorterTest.java @@ -20,7 +20,16 @@ import static org.junit.Assert.fail; +import java.io.IOException; +import java.nio.file.FileVisitResult; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.SimpleFileVisitor; +import java.nio.file.attribute.BasicFileAttributes; + import org.apache.beam.sdk.extensions.sorter.SorterTestUtils.SorterGenerator; +import org.junit.AfterClass; +import org.junit.BeforeClass; import org.junit.Rule; import org.junit.Test; import org.junit.rules.ExpectedException; @@ -31,25 +40,52 @@ @RunWith(JUnit4.class) public class ExternalSorterTest { @Rule public ExpectedException thrown = ExpectedException.none(); + static Path tmpLocation; + + @BeforeClass + public static void setupTempDir() throws IOException { + tmpLocation = Files.createTempDirectory("tmp"); + } + + @AfterClass + public static void cleanupTempDir() throws IOException { + Files.walkFileTree(tmpLocation, new SimpleFileVisitor() { + @Override + public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException { + Files.delete(file); + return FileVisitResult.CONTINUE; + } + + @Override + public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOException { + Files.delete(dir); + return FileVisitResult.CONTINUE; + } + }); + } @Test public void testEmpty() throws Exception { - SorterTestUtils.testEmpty(ExternalSorter.create(new ExternalSorter.Options())); + SorterTestUtils.testEmpty(ExternalSorter.create(new ExternalSorter.Options() + .setTempLocation(tmpLocation.toString()))); } @Test public void testSingleElement() throws Exception { - SorterTestUtils.testSingleElement(ExternalSorter.create(new ExternalSorter.Options())); + SorterTestUtils.testSingleElement(ExternalSorter.create(new ExternalSorter.Options() + .setTempLocation(tmpLocation.toString()))); } @Test public void testEmptyKeyValueElement() throws Exception { - SorterTestUtils.testEmptyKeyValueElement(ExternalSorter.create(new ExternalSorter.Options())); + SorterTestUtils.testEmptyKeyValueElement(ExternalSorter.create(new ExternalSorter.Options() + .setTempLocation(tmpLocation.toString()))); } @Test public void testMultipleIterations() throws Exception { - SorterTestUtils.testMultipleIterations(ExternalSorter.create(new ExternalSorter.Options())); + SorterTestUtils.testMultipleIterations(ExternalSorter.create(new ExternalSorter.Options() + .setTempLocation(tmpLocation.toString()))); } @Test @@ -58,7 +94,8 @@ public void testRandom() throws Exception { new SorterGenerator() { @Override public Sorter generateSorter() throws Exception { - return ExternalSorter.create(new ExternalSorter.Options()); + return ExternalSorter.create(new ExternalSorter.Options() + .setTempLocation(tmpLocation.toString())); } }, 1, @@ -67,13 +104,15 @@ public Sorter generateSorter() throws Exception { @Test public void testAddAfterSort() throws Exception { - SorterTestUtils.testAddAfterSort(ExternalSorter.create(new ExternalSorter.Options()), thrown); + SorterTestUtils.testAddAfterSort(ExternalSorter.create(new ExternalSorter.Options() + .setTempLocation(tmpLocation.toString())), thrown); fail(); } @Test public void testSortTwice() throws Exception { - SorterTestUtils.testSortTwice(ExternalSorter.create(new ExternalSorter.Options()), thrown); + SorterTestUtils.testSortTwice(ExternalSorter.create(new ExternalSorter.Options() + .setTempLocation(tmpLocation.toString())), thrown); fail(); }