diff --git a/src/main/java/io/zonky/test/db/postgres/embedded/EmbeddedPostgres.java b/src/main/java/io/zonky/test/db/postgres/embedded/EmbeddedPostgres.java index 2bd7353a..5714f6a5 100644 --- a/src/main/java/io/zonky/test/db/postgres/embedded/EmbeddedPostgres.java +++ b/src/main/java/io/zonky/test/db/postgres/embedded/EmbeddedPostgres.java @@ -71,6 +71,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Consumer; import java.util.stream.Stream; import static java.nio.file.StandardOpenOption.CREATE; @@ -114,14 +115,14 @@ public class EmbeddedPostgres implements Closeable PgBinaryResolver pgBinaryResolver, ProcessBuilder.Redirect errorRedirector, ProcessBuilder.Redirect outputRedirector) throws IOException { this(parentDirectory, dataDirectory, cleanDataDirectory, postgresConfig, localeConfig, port, connectConfig, - pgBinaryResolver, errorRedirector, outputRedirector, DEFAULT_PG_STARTUP_WAIT, null); + pgBinaryResolver, errorRedirector, outputRedirector, DEFAULT_PG_STARTUP_WAIT, null, null); } EmbeddedPostgres(File parentDirectory, File dataDirectory, boolean cleanDataDirectory, Map postgresConfig, Map localeConfig, int port, Map connectConfig, PgBinaryResolver pgBinaryResolver, ProcessBuilder.Redirect errorRedirector, ProcessBuilder.Redirect outputRedirector, Duration pgStartupWait, - File overrideWorkingDirectory) throws IOException + File overrideWorkingDirectory, Consumer dataDirectoryCustomizer) throws IOException { this.cleanDataDirectory = cleanDataDirectory; this.postgresConfig = new HashMap<>(postgresConfig); @@ -158,6 +159,11 @@ public class EmbeddedPostgres implements Closeable } lock(); + + if (dataDirectoryCustomizer != null) { + dataDirectoryCustomizer.accept(dataDirectory); + } + startPostmaster(); } @@ -495,6 +501,7 @@ public static class Builder private final Map connectConfig = new HashMap<>(); private PgBinaryResolver pgBinaryResolver = DefaultPostgresBinaryResolver.INSTANCE; private Duration pgStartupWait = DEFAULT_PG_STARTUP_WAIT; + private Consumer dataDirectoryCustomizer; private ProcessBuilder.Redirect errRedirector = ProcessBuilder.Redirect.PIPE; private ProcessBuilder.Redirect outRedirector = ProcessBuilder.Redirect.PIPE; @@ -573,6 +580,11 @@ public Builder setPgBinaryResolver(PgBinaryResolver pgBinaryResolver) { return this; } + public Builder setDataDirectoryCustomizer(final Consumer dataDirectoryCustomizer) { + this.dataDirectoryCustomizer = dataDirectoryCustomizer; + return this; + } + public EmbeddedPostgres start() throws IOException { if (builderPort == 0) { @@ -583,7 +595,7 @@ public EmbeddedPostgres start() throws IOException { } return new EmbeddedPostgres(parentDirectory, builderDataDirectory, builderCleanDataDirectory, config, localeConfig, builderPort, connectConfig, pgBinaryResolver, errRedirector, outRedirector, - pgStartupWait, overrideWorkingDirectory); + pgStartupWait, overrideWorkingDirectory, dataDirectoryCustomizer); } @Override diff --git a/src/test/java/io/zonky/test/db/postgres/embedded/EmbeddedPostgresTest.java b/src/test/java/io/zonky/test/db/postgres/embedded/EmbeddedPostgresTest.java index b752dc5b..75f587dd 100644 --- a/src/test/java/io/zonky/test/db/postgres/embedded/EmbeddedPostgresTest.java +++ b/src/test/java/io/zonky/test/db/postgres/embedded/EmbeddedPostgresTest.java @@ -22,6 +22,7 @@ import java.sql.Connection; import java.sql.ResultSet; import java.sql.Statement; +import java.util.concurrent.atomic.AtomicBoolean; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; @@ -47,10 +48,18 @@ public void testEmbeddedPg() throws Exception @Test public void testEmbeddedPgCreationWithNestedDataDirectory() throws Exception { + AtomicBoolean called = new AtomicBoolean(false); + Path dataDir = Files.createDirectories(tf.resolve("data-dir-parent").resolve("data-dir")); try (EmbeddedPostgres pg = EmbeddedPostgres.builder() - .setDataDirectory(Files.createDirectories(tf.resolve("data-dir-parent").resolve("data-dir"))) + .setDataDirectory(dataDir) + .setDataDirectoryCustomizer(dd -> { + called.set(true); + assertEquals(dataDir, dd.toPath()); + assertTrue(Files.isRegularFile(dd.toPath().resolve("pg_hba.conf"))); + }) .start()) { // nothing to do } + assertTrue(called.get()); } }