/** * This file is part of Graylog. * * Graylog is free software: you can redistribute it and/or modify * it under the terms of the GNU General Public License as published by * the Free Software Foundation, either version 3 of the License, or * (at your option) any later version. * * Graylog is distributed in the hope that it will be useful, * but WITHOUT ANY WARRANTY; without even the implied warranty of * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the * GNU General Public License for more details. * * You should have received a copy of the GNU General Public License * along with Graylog. If not, see <http://www.gnu.org/licenses/>. */ package integration; import integration.util.graylog.ServerHelper; import integration.util.mongodb.MongodbSeed; import org.junit.rules.MethodRule; import org.junit.runners.model.FrameworkMethod; import org.junit.runners.model.Statement; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; import java.net.MalformedURLException; import java.net.URISyntaxException; import java.net.URL; import java.net.UnknownHostException; import static com.google.common.base.Strings.emptyToNull; import static org.apache.commons.lang3.ObjectUtils.firstNonNull; public class MongoDbSeedRule implements MethodRule { private static final Logger log = LoggerFactory.getLogger(MongoDbSeedRule.class); private static final String CK_RESOURCEPREFIX = "integration/seeds/mongodb/"; private static final String CK_JSONSUFFIX = ".json"; @Override public Statement apply(Statement base, FrameworkMethod method, Object target) { final MongoDbSeed annotation = firstNonNull(method.getAnnotation(MongoDbSeed.class), method.getDeclaringClass().getAnnotation(MongoDbSeed.class)); if (annotation != null) { final String databaseName = firstNonNull(emptyToNull(annotation.database()), IntegrationTestsConfig.getMongodbDatabase()); final MongodbSeed mongodbSeed; try { mongodbSeed = new MongodbSeed(databaseName); } catch (UnknownHostException e) { final String msg = "Unable to seed database: "; log.error(msg, e); return new IgnoreStatement(msg + e.toString()); } final ServerHelper graylogController = new ServerHelper(); final String nodeId; try { nodeId = graylogController.getNodeId(); } catch (MalformedURLException | URISyntaxException e) { final String msg = "Unable to determine Graylog node id for seeding: "; log.error(msg, e); return new IgnoreStatement(msg + e.toString()); } if (annotation.locations().length > 0) { for (String location : annotation.locations()) { final URL seedUrl = findFirstInSearchPath(location, method.getDeclaringClass()); if (seedUrl == null) { throw new RuntimeException("Unable to find seed data " + location + " for " + method.toString()); } try { log.debug("Using seed data from {}", seedUrl); mongodbSeed.loadDataset(seedUrl, nodeId); } catch (IOException | URISyntaxException e) { throw new RuntimeException("Unable to read seed data: ", e); } } } else { final URL seedUrl = findFirstInSearchPath(method.getName(), method.getDeclaringClass()); try { if (seedUrl != null) { log.debug("Using seed data from {}", seedUrl); mongodbSeed.loadDataset(seedUrl, nodeId); } } catch (IOException | URISyntaxException e) { log.debug("MongoDB seed annotation present, but neither explicit location passed, nor inferenced location available, just cleaning databse. Exception was: ", e); } } } return base; } private URL findFirstInSearchPath(String suffix, Class<?> testClass) { return firstNonNull( firstNonNull(getClassSpecificResource(testClass, suffix), getClassSpecificResource(testClass, suffix + CK_JSONSUFFIX)), firstNonNull(getGlobalResource(suffix), getGlobalResource(suffix + CK_JSONSUFFIX)) ); } private URL getGlobalResource(String name) { return Thread.currentThread().getContextClassLoader().getResource(CK_RESOURCEPREFIX + name); } private URL getClassSpecificResource(Class klazz, String name) { return klazz.getResource(klazz.getSimpleName() + "/" + name); } }