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..2af7c9aa 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 @@ -101,6 +101,7 @@ public class EmbeddedPostgres implements Closeable private final Map postgresConfig; private final Map localeConfig; private final Map connectConfig; + private final Map environmentVariables; private volatile FileOutputStream lockStream; private volatile FileLock lock; @@ -108,17 +109,19 @@ public class EmbeddedPostgres implements Closeable private final ProcessBuilder.Redirect errorRedirector; private final ProcessBuilder.Redirect outputRedirector; + private Process process; EmbeddedPostgres(File parentDirectory, File dataDirectory, boolean cleanDataDirectory, Map postgresConfig, Map localeConfig, int port, Map connectConfig, - PgBinaryResolver pgBinaryResolver, ProcessBuilder.Redirect errorRedirector, ProcessBuilder.Redirect outputRedirector) throws IOException + Map environmentVariables, PgBinaryResolver pgBinaryResolver, ProcessBuilder.Redirect errorRedirector, ProcessBuilder.Redirect outputRedirector) throws IOException { - this(parentDirectory, dataDirectory, cleanDataDirectory, postgresConfig, localeConfig, port, connectConfig, + this(parentDirectory, dataDirectory, cleanDataDirectory, postgresConfig, localeConfig, port, connectConfig, environmentVariables, pgBinaryResolver, errorRedirector, outputRedirector, DEFAULT_PG_STARTUP_WAIT, null); } EmbeddedPostgres(File parentDirectory, File dataDirectory, boolean cleanDataDirectory, Map postgresConfig, Map localeConfig, int port, Map connectConfig, + Map environmentVariables, PgBinaryResolver pgBinaryResolver, ProcessBuilder.Redirect errorRedirector, ProcessBuilder.Redirect outputRedirector, Duration pgStartupWait, File overrideWorkingDirectory) throws IOException @@ -127,6 +130,7 @@ public class EmbeddedPostgres implements Closeable this.postgresConfig = new HashMap<>(postgresConfig); this.localeConfig = new HashMap<>(localeConfig); this.connectConfig = new HashMap<>(connectConfig); + this.environmentVariables = new HashMap<>(environmentVariables); this.port = port; this.pgDir = prepareBinaries(pgBinaryResolver, overrideWorkingDirectory); this.errorRedirector = errorRedirector; @@ -158,7 +162,11 @@ public class EmbeddedPostgres implements Closeable } lock(); - startPostmaster(); + this.process = startPostmaster(); + } + + public Process getProcess() { + return this.process; } public DataSource getTemplateDatabase() @@ -251,7 +259,7 @@ private void initdb() LOG.info("{} initdb completed in {}", instanceId, watch); } - private void startPostmaster() throws IOException + private Process startPostmaster() throws IOException { final StopWatch watch = new StopWatch(); watch.start(); @@ -283,6 +291,7 @@ private void startPostmaster() throws IOException Runtime.getRuntime().addShutdownHook(newCloserThread()); waitForServerStartup(watch); + return postmaster; } private List createInitOptions() @@ -490,6 +499,7 @@ public static class Builder private File builderDataDirectory; private final Map config = new HashMap<>(); private final Map localeConfig = new HashMap<>(); + private final Map environmentVariables = new HashMap<>(); private boolean builderCleanDataDirectory = true; private int builderPort = 0; private final Map connectConfig = new HashMap<>(); @@ -546,6 +556,11 @@ public Builder setConnectConfig(String key, String value) { connectConfig.put(key, value); return this; } + + public Builder setEnvironmentVariable(String key, String value) { + environmentVariables.put(key, value); + return this; + } public Builder setOverrideWorkingDirectory(File workingDirectory) { overrideWorkingDirectory = workingDirectory; @@ -582,7 +597,7 @@ public EmbeddedPostgres start() throws IOException { builderDataDirectory = Files.createTempDirectory("epg").toFile(); } return new EmbeddedPostgres(parentDirectory, builderDataDirectory, builderCleanDataDirectory, config, - localeConfig, builderPort, connectConfig, pgBinaryResolver, errRedirector, outRedirector, + localeConfig, builderPort, connectConfig, environmentVariables, pgBinaryResolver, errRedirector, outRedirector, pgStartupWait, overrideWorkingDirectory); } @@ -602,6 +617,7 @@ public boolean equals(Object o) { Objects.equals(config, builder.config) && Objects.equals(localeConfig, builder.localeConfig) && Objects.equals(connectConfig, builder.connectConfig) && + Objects.equals(environmentVariables, builder.environmentVariables) && Objects.equals(pgBinaryResolver, builder.pgBinaryResolver) && Objects.equals(pgStartupWait, builder.pgStartupWait) && Objects.equals(errRedirector, builder.errRedirector) && @@ -610,7 +626,7 @@ public boolean equals(Object o) { @Override public int hashCode() { - return Objects.hash(parentDirectory, builderDataDirectory, config, localeConfig, builderCleanDataDirectory, builderPort, connectConfig, pgBinaryResolver, pgStartupWait, errRedirector, outRedirector); + return Objects.hash(parentDirectory, builderDataDirectory, config, localeConfig, builderCleanDataDirectory, builderPort, connectConfig, environmentVariables, pgBinaryResolver, pgStartupWait, errRedirector, outRedirector); } } @@ -894,6 +910,9 @@ public void applyTo(ProcessBuilder builder, List arguments) { command.addAll(arguments); builder.command(command); + + builder.environment().clear(); + builder.environment().putAll(environmentVariables); } } }