001/*
002 * Licensed under the Apache License, Version 2.0 (the "License");
003 * you may not use this file except in compliance with the License.
004 * You may obtain a copy of the License at
005 *
006 * http://www.apache.org/licenses/LICENSE-2.0
007 *
008 * Unless required by applicable law or agreed to in writing, software
009 * distributed under the License is distributed on an "AS IS" BASIS,
010 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
011 * See the License for the specific language governing permissions and
012 * limitations under the License.
013 */
014
015package de.softwareforge.testing.postgres.junit5;
016
017import static com.google.common.base.Preconditions.checkNotNull;
018import static com.google.common.base.Preconditions.checkState;
019
020import de.softwareforge.testing.postgres.embedded.DatabaseInfo;
021import de.softwareforge.testing.postgres.embedded.DatabaseManager;
022import de.softwareforge.testing.postgres.embedded.DatabaseManager.DatabaseManagerBuilder;
023import de.softwareforge.testing.postgres.embedded.EmbeddedPostgres;
024
025import jakarta.annotation.Nonnull;
026import java.lang.reflect.Type;
027import java.sql.SQLException;
028import java.util.UUID;
029import javax.sql.DataSource;
030
031import com.google.common.annotations.VisibleForTesting;
032import org.junit.jupiter.api.extension.AfterAllCallback;
033import org.junit.jupiter.api.extension.AfterEachCallback;
034import org.junit.jupiter.api.extension.BeforeAllCallback;
035import org.junit.jupiter.api.extension.BeforeEachCallback;
036import org.junit.jupiter.api.extension.ExtensionContext;
037import org.junit.jupiter.api.extension.ExtensionContext.Namespace;
038import org.junit.jupiter.api.extension.ExtensionContext.Store;
039import org.junit.jupiter.api.extension.ParameterContext;
040import org.junit.jupiter.api.extension.ParameterResolutionException;
041import org.junit.jupiter.api.extension.ParameterResolver;
042import org.slf4j.Logger;
043import org.slf4j.LoggerFactory;
044
045/**
046 * A <a href="https://junit.org/junit5/docs/current/user-guide/#extensions">JUnit 5 extension</a> that manages an embedded PostgreSQL database server.
047 * <p>
048 * This extension can provide the {@link EmbeddedPostgres} instance, a {@link DatabaseInfo} or a {@link DataSource} object as test parameters.
049 *
050 * <ul>
051 * <li>Using a {@link DatabaseInfo} parameter is equivalent to calling {@link EmbeddedPgExtension#createDatabaseInfo()}.</li>
052 * <li>Using a {@link DataSource} parameter is equivalent to calling {@link EmbeddedPgExtension#createDataSource()}.</li>
053 * </ul>
054 */
055public final class EmbeddedPgExtension implements BeforeAllCallback, AfterAllCallback, BeforeEachCallback, AfterEachCallback, ParameterResolver {
056
057    private static final Logger LOG = LoggerFactory.getLogger(EmbeddedPgExtension.class);
058
059    // multiple instances must use different namespaces
060    private final Namespace pgNamespace = Namespace.create(UUID.randomUUID());
061
062    private final DatabaseManager.Builder<DatabaseManager> databaseManagerBuilder;
063
064    private volatile DatabaseManager databaseManager = null;
065
066    private EmbeddedPgExtension(DatabaseManager.Builder<DatabaseManager> databaseManagerBuilder) {
067        this.databaseManagerBuilder = databaseManagerBuilder;
068    }
069
070    /**
071     * Creates a new {@link EmbeddedPgExtensionBuilder} that allows further customization of the {@link EmbeddedPgExtension}. The resulting extension manages
072     * the database server in multi-mode (creating multiple databases).
073     *
074     * @return A {@link EmbeddedPgExtensionBuilder} instance. Never null.
075     */
076    @Nonnull
077    static EmbeddedPgExtensionBuilder multiDatabase() {
078        return new EmbeddedPgExtensionBuilder(true);
079    }
080
081    /**
082     * Creates a new {@link EmbeddedPgExtensionBuilder} that allows further customization of the {@link EmbeddedPgExtension}. The resulting extension manages
083     * the database server in single-mode (using a single database instance).
084     *
085     * @return A {@link EmbeddedPgExtensionBuilder} instance. Never null.
086     */
087    @Nonnull
088    static EmbeddedPgExtensionBuilder singleDatabase() {
089        return new EmbeddedPgExtensionBuilder(false);
090    }
091
092    /**
093     * Default constructor which allows using this extension with the {@link org.junit.jupiter.api.extension.ExtendWith} annotation.
094     * <p>
095     * This is equivalent to using <pre>{@code
096     *     @RegisterExtension
097     *     public static EmbeddedPgExtension pg = MultiDatabaseBuilder.instanceWithDefaults().build();
098     *     }</pre>
099     *
100     * @since 3.0
101     */
102    public EmbeddedPgExtension() {
103        this(new DatabaseManagerBuilder(true).withInstancePreparer(EmbeddedPostgres.Builder::withDefaults));
104    }
105
106    /**
107     * Returns a data source. Depending on the mode it returns a datasource connected to the same database (single mode) or a new database (multi mode) on every
108     * call.
109     *
110     * @return A {@link DataSource} instance. This is never null.
111     * @throws SQLException If a problem connecting to the database occurs.
112     */
113    @Nonnull
114    public DataSource createDataSource() throws SQLException {
115        return createDatabaseInfo().asDataSource();
116    }
117
118    @VisibleForTesting
119    EmbeddedPostgres getEmbeddedPostgres() {
120        return databaseManager.getEmbeddedPostgres();
121    }
122
123    /**
124     * Returns a new {@link DatabaseInfo} describing a database connection.
125     * <p>
126     * Depending on the mode, this either describes the same database (single mode) or a new database (multi mode).
127     *
128     * @return A {@link DatabaseInfo} instance. This is never null.
129     * @throws SQLException If a problem connecting to the database occurs.
130     */
131    @Nonnull
132    public DatabaseInfo createDatabaseInfo() throws SQLException {
133        checkState(databaseManager != null, "no before method has been called!");
134
135        DatabaseInfo databaseInfo = databaseManager.getDatabaseInfo();
136        LOG.info("Connection to {}", databaseInfo.asJdbcUrl());
137        return databaseInfo;
138    }
139
140    @Override
141    public void beforeAll(@Nonnull ExtensionContext extensionContext) throws Exception {
142        checkNotNull(extensionContext, "extensionContext is null");
143
144        Store pgStore = extensionContext.getStore(pgNamespace);
145
146        TestingContext testingContext = pgStore.getOrComputeIfAbsent(TestingContext.TESTING_CONTEXT_KEY,
147                k -> new TestingContext(extensionContext.getUniqueId(), databaseManagerBuilder.build()),
148                TestingContext.class);
149
150        this.databaseManager = testingContext.start(extensionContext.getUniqueId());
151    }
152
153    @Override
154    public void afterAll(@Nonnull ExtensionContext extensionContext) throws Exception {
155        checkNotNull(extensionContext, "extensionContext is null");
156
157        Store pgStore = extensionContext.getStore(pgNamespace);
158        TestingContext testingContext = pgStore.get(TestingContext.TESTING_CONTEXT_KEY, TestingContext.class);
159
160        if (testingContext != null) {
161            this.databaseManager = testingContext.stop(extensionContext.getUniqueId());
162        }
163    }
164
165    @Override
166    public void beforeEach(@Nonnull ExtensionContext extensionContext) throws Exception {
167        checkNotNull(extensionContext, "extensionContext is null");
168
169        Store pgStore = extensionContext.getStore(pgNamespace);
170        TestingContext testingContext = pgStore.getOrComputeIfAbsent(TestingContext.TESTING_CONTEXT_KEY,
171                k -> new TestingContext(extensionContext.getUniqueId(), databaseManagerBuilder.build()),
172                TestingContext.class);
173
174        this.databaseManager = testingContext.start(extensionContext.getUniqueId());
175    }
176
177    @Override
178    public void afterEach(@Nonnull ExtensionContext extensionContext) throws Exception {
179        checkNotNull(extensionContext, "extensionContext is null");
180
181        Store pgStore = extensionContext.getStore(pgNamespace);
182        TestingContext testingContext = pgStore.get(TestingContext.TESTING_CONTEXT_KEY, TestingContext.class);
183
184        if (testingContext != null) {
185            this.databaseManager = testingContext.stop(extensionContext.getUniqueId());
186        }
187    }
188
189    @Override
190    public boolean supportsParameter(@Nonnull ParameterContext parameterContext, ExtensionContext extensionContext) {
191        Type type = parameterContext.getParameter().getType();
192        return type == EmbeddedPostgres.class
193                || type == DatabaseInfo.class
194                || type == DataSource.class;
195    }
196
197    @Override
198    public Object resolveParameter(@Nonnull ParameterContext parameterContext, ExtensionContext extensionContext) {
199        Type type = parameterContext.getParameter().getType();
200        try {
201            if (type == EmbeddedPostgres.class) {
202                return getEmbeddedPostgres();
203            } else if (type == DatabaseInfo.class) {
204                return createDatabaseInfo();
205
206            } else if (type == DataSource.class) {
207                return createDataSource();
208            }
209        } catch (SQLException e) {
210            throw new ParameterResolutionException("Could not create " + type.getTypeName() + " instance", e);
211        }
212        return null;
213    }
214
215    /**
216     * Builder for {@link EmbeddedPgExtension} customization.
217     */
218    public static final class EmbeddedPgExtensionBuilder extends DatabaseManager.Builder<EmbeddedPgExtension> {
219
220        private EmbeddedPgExtensionBuilder(boolean multiMode) {
221            super(multiMode);
222        }
223
224        /**
225         * Create a {@link EmbeddedPgExtension} instance.
226         *
227         * @return A {@link EmbeddedPgExtension} instance. Is never null.
228         */
229        @Override
230        @Nonnull
231        public EmbeddedPgExtension build() {
232            DatabaseManager.Builder<DatabaseManager> databaseManagerBuilder = new DatabaseManagerBuilder(multiMode)
233                    .withDatabasePreparers(databasePreparers.build())
234                    .withInstancePreparers(instancePreparers.build());
235
236            return new EmbeddedPgExtension(databaseManagerBuilder);
237        }
238    }
239
240    private static final class TestingContext {
241
242        private static final Object TESTING_CONTEXT_KEY = new Object();
243
244        private final String id;
245        private final DatabaseManager databaseManager;
246
247        private TestingContext(String id, DatabaseManager databaseManager) {
248            this.id = id;
249            this.databaseManager = databaseManager;
250        }
251
252        private DatabaseManager start(String id) throws Exception {
253            if (this.id.equals(id)) {
254                databaseManager.start();
255            }
256
257            return databaseManager;
258        }
259
260        private DatabaseManager stop(String id) throws Exception {
261            if (this.id.equals(id)) {
262                databaseManager.close();
263                return null;
264            }
265
266            return databaseManager;
267        }
268    }
269}