org.ambraproject.testutils.EmbeddedSolrServerFactory.java Source code

Java tutorial

Introduction

Here is the source code for org.ambraproject.testutils.EmbeddedSolrServerFactory.java

Source

/*
 * $HeadURL$
 * $Id$
 * Copyright (c) 2006-2012 by Public Library of Science http://plos.org http://ambraproject.org
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.ambraproject.testutils;

import org.ambraproject.service.search.SolrServerFactory;
import org.apache.commons.io.FileUtils;
import org.apache.commons.io.IOUtils;
import org.apache.solr.client.solrj.SolrServer;
import org.apache.solr.client.solrj.embedded.EmbeddedSolrServer;
import org.apache.solr.client.solrj.response.UpdateResponse;
import org.apache.solr.common.SolrInputDocument;
import org.apache.solr.core.CoreContainer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.Map;

/**
 * Class added for unit tests to use an in-memory embedded solr server.
 * <p/>
 * Extends the Ambra SolrServerFactory so it can be passed to existing beans. Documents are added to the server
 * by passing a map to the addDocument method.
 *
 * @author Alex Kudlick Date: 5/13/11
 *         <p/>
 *         org.ambraproject.testutils
 */
public class EmbeddedSolrServerFactory extends SolrServerFactory {

    private static Logger log = LoggerFactory.getLogger(EmbeddedSolrServerFactory.class);

    private SolrServer server;

    private String solrHome;

    private String destSolrXmlFile;

    /**
     * Constructor.
     *
     * TODO: this constructor does a lot of IO, and can throw IOException, not ideal
     * in a constructor.  Refactor to eliminate this.
     */
    public EmbeddedSolrServerFactory() throws IOException {

        // Build a directory structure for the embedded solr server's solr home,
        // and copy configuration files there from the classpath.
        solrHome = System.getProperty("java.io.tmpdir") + File.separator + "EmbeddedSolrTest_"
                + System.currentTimeMillis();
        File conf = new File(solrHome + File.separator + "collection1" + File.separator + "conf");
        if (!conf.mkdirs()) {
            throw new RuntimeException("Could not create dir " + conf.getCanonicalPath());
        }
        destSolrXmlFile = solrHome + File.separator + "solr.xml";
        copyResource("solr/collection1/conf/test-solr.xml", destSolrXmlFile);
        copyResource("solr/collection1/conf/test-solr-config.xml",
                conf.getCanonicalPath() + File.separator + "solrconfig.xml");
        copyResource("solr/collection1/conf/test-solr-schema.xml",
                conf.getCanonicalPath() + File.separator + "schema.xml");
        copyResource("solr/collection1/conf/stopwords.txt",
                conf.getCanonicalPath() + File.separator + "stopwords.txt");
        copyResource("solr/collection1/conf/author_stopwords.txt",
                conf.getCanonicalPath() + File.separator + "author_stopwords.txt");

        System.setProperty("solr.solr.home", solrHome);
        CoreContainer coreContainer = new CoreContainer(solrHome, new File(destSolrXmlFile));
        server = new EmbeddedSolrServer(coreContainer, "collection1");
        log.info("EmbeddedSolrServer started with solr home " + solrHome);
    }

    /**
     * Copies a file from the classpath to the given destination.
     *
     * @param resource resource location relative to the classpath
     * @param destination destination file path
     * @throws IOException
     */
    private void copyResource(String resource, String destination) throws IOException {
        IOUtils.copy(getClass().getClassLoader().getResourceAsStream(resource),
                new FileOutputStream(new File(destination)));
    }

    public void tearDown() throws Exception {
        server = null;
        FileUtils.deleteQuietly(new File(solrHome));
    }

    /**
     * Add a document to the server stored here. Note that "id" is a required field
     *
     * @param document - a map from solr field names to values.  The value array should have more than one entry only for
     *                 fields that are multivalued (see the test schema.xml)
     * @throws Exception - from the server.add() method
     */
    public void addDocument(Map<String, String[]> document) throws Exception {
        SolrInputDocument inputDocument = new SolrInputDocument();
        for (String fieldName : document.keySet()) {
            for (String value : document.get(fieldName)) {
                inputDocument.addField(fieldName, value);
            }
        }
        server.add(inputDocument);
        server.commit();
    }

    /**
     * Add a document to the server stored here.  Note that "id" is a required field
     *
     * @param document each row should be a field, with the first entry the field name, and the rest of the entries the values for the field.
     *                 Only multi-valued fields should have more than one value
     * @throws Exception from there server.add() method
     */
    public void addDocument(String[][] document) throws Exception {
        SolrInputDocument inputDocument = new SolrInputDocument();
        for (String[] row : document) {
            String fieldName = row[0];
            for (String value : Arrays.copyOfRange(row, 1, row.length)) {
                inputDocument.addField(fieldName, value);
            }
        }
        server.add(inputDocument);
        server.commit();
    }

    /**
     * Delete all documents in solr
     */
    public void deleteAllDocuments() {
        try {
            UpdateResponse updateResponse = server.deleteByQuery("*:*");
        } catch (Exception e) {
            log.error("Failed to delete all documents in solr.", e);
        }
    }

    /**
     * Get Embedded Solr Server instance
     *
     * @return Solr server
     */
    @Override
    public SolrServer getServer() {
        return server;
    }
}