<commit_before>package org.mongodb.morphia.aggregation;

import com.mongodb.DBCollection;
import org.bson.types.ObjectId;
import org.junit.Assert;
import org.junit.Assume;
import org.junit.Test;
import org.mongodb.morphia.TestBase;
import org.mongodb.morphia.annotations.AlsoLoad;
import org.mongodb.morphia.annotations.Embedded;
import org.mongodb.morphia.annotations.Entity;
import org.mongodb.morphia.annotations.Id;
import org.mongodb.morphia.annotations.Property;
import org.mongodb.morphia.logging.Logger;
import org.mongodb.morphia.logging.MorphiaLoggerFactory;
import org.mongodb.morphia.query.MorphiaIterator;
import org.mongodb.morphia.query.Query;
import org.zeroturnaround.exec.ProcessExecutor;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.TimeoutException;

import static java.lang.String.format;
import static org.mongodb.morphia.aggregation.Group.average;
import static org.mongodb.morphia.aggregation.Group.first;
import static org.mongodb.morphia.aggregation.Group.grouping;
import static org.mongodb.morphia.aggregation.Group.id;
import static org.mongodb.morphia.aggregation.Group.last;
import static org.mongodb.morphia.aggregation.Group.sum;
import static org.mongodb.morphia.aggregation.Projection.projection;
import static org.mongodb.morphia.aggregation.Sort.ascending;

/**
 * These tests recreate the example zip code data set aggregations as found in the official documentation.
 *
 * @see <a href="http://docs.mongodb.org/manual/tutorial/aggregation-zip-code-data-set/">Aggregation with the Zip Code Data Set</a>
 */
public class ZipCodeDataSetTest extends TestBase {
    public static final String MONGO_IMPORT;
    private static final Logger LOG = MorphiaLoggerFactory.get(ZipCodeDataSetTest.class);

    static {
        String property = System.getProperty("mongodb_server");
        String serverType = property != null ? property.replaceAll("-release", "") : "UNKNOWN";
        String path = format("/mnt/jenkins/mongodb/%s/%s/bin/mongoimport", serverType, property);
        if (new File(path).exists()) {
            MONGO_IMPORT = path;
        } else {
            MONGO_IMPORT = "/usr/local/bin/mongoimport";
        }
    }

    public void installSampleData() throws IOException, TimeoutException, InterruptedException {
        File file = new File("zips.json");
        if (!file.exists()) {
            file = new File(System.getProperty("java.io.tmpdir"), "zips.json");
            if (!file.exists()) {
                download(new URL("http://media.mongodb.org/zips.json"), file);
            }
        }
        DBCollection zips = getDb().getCollection("zips");
        if (zips.count() == 0) {
            new ProcessExecutor().command(MONGO_IMPORT,
                                          "--db", getDb().getName(),
                                          "--collection", "zipcodes",
                                          "--file", file.getAbsolutePath())
                                 .redirectError(System.err)
                                 .execute();
        }
    }

    private void download(final URL url, final File file) throws IOException {
        LOG.info("Downloading zip data set to " + file);
        InputStream inputStream = url.openStream();
        FileOutputStream outputStream = new FileOutputStream(file);
        try {
            byte[] read = new byte[49152];
            int count;
            while ((count = inputStream.read(read)) != -1) {
                outputStream.write(read, 0, count);
            }
        } finally {
            inputStream.close();
            outputStream.close();
        }
    }

    @Test
    public void populationsAbove10M() throws IOException, TimeoutException, InterruptedException {
        Assume.assumeTrue(new File(MONGO_IMPORT).exists());
        installSampleData();
        Query<Object> query = getDs().getQueryFactory().createQuery(getDs());

        AggregationPipeline pipeline
            = getDs().<City, Population>createAggregation(City.class)
                     .group("state", grouping("totalPop", sum("pop")))
                     .match(query.field("totalPop").greaterThanOrEq(10000000));


        validate(pipeline.aggregate(Population.class), "CA", 29754890);
        validate(pipeline.aggregate(Population.class), "OH", 10846517);
    }

    @Test
    public void averageCitySizeByState() throws InterruptedException, TimeoutException, IOException {
        Assume.assumeTrue(new File(MONGO_IMPORT).exists());
        installSampleData();
        AggregationPipeline pipeline = getDs().<City, Population>createAggregation(City.class)
                                              .group(id(grouping("state"), grouping("city")), grouping("pop", sum("pop")))
                                              .group("_id.state", grouping("avgCityPop", average("pop")));
        validate(pipeline.aggregate(Population.class), "MN", 5372);
    }

    @Test
    public void smallestAndLargestCities() throws InterruptedException, TimeoutException, IOException {
        Assume.assumeTrue(new File(MONGO_IMPORT).exists());
        installSampleData();
        getMorphia().mapPackage(getClass().getPackage().getName());
        AggregationPipeline pipeline = getDs().<City, State>createAggregation(City.class)

                                              .group(id(grouping("state"), grouping("city")), grouping("pop", sum("pop")))

                                              .sort(ascending("pop"))

                                              .group("_id.state",
                                                     grouping("biggestCity", last("_id.city")),
                                                     grouping("biggestPop", last("pop")),
                                                     grouping("smallestCity", first("_id.city")),
                                                     grouping("smallestPop", first("pop"))
                                                    )

                                              .project(projection("_id").suppress(),
                                                       projection("state", "_id"),
                                                       projection("biggestCity",
                                                                  projection("name", "biggestCity"),
                                                                  projection("pop", "biggestPop")
                                                                 ),
                                                       projection("smallestCity",
                                                                  projection("name", "smallestCity"),
                                                                  projection("pop", "smallestPop")
                                                                 )
                                                      );

        Iterator<State> iterator = pipeline.aggregate(State.class);
        try {
            Map<String, State> states = new HashMap<String, State>();
            while (iterator.hasNext()) {
                State state = iterator.next();<commit_msg>fixing the correct file just for grins.<commit_after>package org.mongodb.morphia.aggregation;

import com.mongodb.DBCollection;
import org.bson.types.ObjectId;
import org.junit.Assert;
import org.junit.Assume;
import org.junit.Test;
import org.mongodb.morphia.TestBase;
import org.mongodb.morphia.annotations.AlsoLoad;
import org.mongodb.morphia.annotations.Embedded;
import org.mongodb.morphia.annotations.Entity;
import org.mongodb.morphia.annotations.Id;
import org.mongodb.morphia.annotations.Property;
import org.mongodb.morphia.logging.Logger;
import org.mongodb.morphia.logging.MorphiaLoggerFactory;
import org.mongodb.morphia.query.MorphiaIterator;
import org.mongodb.morphia.query.Query;
import org.zeroturnaround.exec.ProcessExecutor;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.TimeoutException;

import static java.lang.String.format;
import static org.mongodb.morphia.aggregation.Group.average;
import static org.mongodb.morphia.aggregation.Group.first;
import static org.mongodb.morphia.aggregation.Group.grouping;
import static org.mongodb.morphia.aggregation.Group.id;
import static org.mongodb.morphia.aggregation.Group.last;
import static org.mongodb.morphia.aggregation.Group.sum;
import static org.mongodb.morphia.aggregation.Projection.projection;
import static org.mongodb.morphia.aggregation.Sort.ascending;

/**
 * These tests recreate the example zip code data set aggregations as found in the official documentation.
 *
 * @see <a href="http://docs.mongodb.org/manual/tutorial/aggregation-zip-code-data-set/">Aggregation with the Zip Code Data Set</a>
 */
public class ZipCodeDataSetTest extends TestBase {
    public static final String MONGO_IMPORT;
    private static final Logger LOG = MorphiaLoggerFactory.get(ZipCodeDataSetTest.class);

    static {
        String property = System.getProperty("mongodb_server");
        String serverType = property != null ? property.replaceAll("-release", "") : "UNKNOWN";
        String path = format("/mnt/jenkins/mongodb/%s/%s/bin/mongoimport", serverType, property);
        if (new File(path).exists()) {
            MONGO_IMPORT = path;
        } else {
            MONGO_IMPORT = "/usr/local/bin/mongoimport";
        }
    }

    public void installSampleData() throws IOException, TimeoutException, InterruptedException {
        File file = new File("zips.json");
        if (!file.exists()) {
            file = new File(System.getProperty("java.io.tmpdir"), "zips.json");
            if (!file.exists()) {
                download(new URL("http://media.mongodb.org/zips.json"), file);
            }
        }
        DBCollection zips = getDb().getCollection("zips");
        if (zips.count() == 0) {
            new ProcessExecutor().command(MONGO_IMPORT,
                                          "--db", getDb().getName(),
                                          "--collection", "zipcodes",
                                          "--file", file.getAbsolutePath())
                                 .redirectError(System.err)
                                 .execute();
        }
    }

    private void download(final URL url, final File file) throws IOException {
        LOG.info("Downloading zip data set to " + file);
        InputStream inputStream = url.openStream();
        FileOutputStream outputStream = new FileOutputStream(file);
        try {
            byte[] read = new byte[49152];
            int count;
            while ((count = inputStream.read(read)) != -1) {
                outputStream.write(read, 0, count);
            }
        } finally {
            inputStream.close();
            outputStream.close();
        }
    }

    @Test
    public void populationsAbove10M() throws IOException, TimeoutException, InterruptedException {
        Assume.assumeTrue(new File(MONGO_IMPORT).exists());
        installSampleData();
        Query<Object> query = getDs().getQueryFactory().createQuery(getDs());

        AggregationPipeline pipeline
            = getDs().createAggregation(City.class)
                     .group("state", grouping("totalPop", sum("pop")))
                     .match(query.field("totalPop").greaterThanOrEq(10000000));


        validate(pipeline.aggregate(Population.class), "CA", 29754890);
        validate(pipeline.aggregate(Population.class), "OH", 10846517);
    }

    @Test
    public void averageCitySizeByState() throws InterruptedException, TimeoutException, IOException {
        Assume.assumeTrue(new File(MONGO_IMPORT).exists());
        installSampleData();
        AggregationPipeline pipeline = getDs().createAggregation(City.class)
                                              .group(id(grouping("state"), grouping("city")), grouping("pop", sum("pop")))
                                              .group("_id.state", grouping("avgCityPop", average("pop")));
        validate(pipeline.aggregate(Population.class), "MN", 5372);
    }

    @Test
    public void smallestAndLargestCities() throws InterruptedException, TimeoutException, IOException {
        Assume.assumeTrue(new File(MONGO_IMPORT).exists());
        installSampleData();
        getMorphia().mapPackage(getClass().getPackage().getName());
        AggregationPipeline pipeline = getDs().createAggregation(City.class)

                                              .group(id(grouping("state"), grouping("city")), grouping("pop", sum("pop")))

                                              .sort(ascending("pop"))

                                              .group("_id.state",
                                                     grouping("biggestCity", last("_id.city")),
                                                     grouping("biggestPop", last("pop")),
                                                     grouping("smallestCity", first("_id.city")),
                                                     grouping("smallestPop", first("pop")))

                                              .project(projection("_id").suppress(),
                                                       projection("state", "_id"),
                                                       projection("biggestCity",
                                                                  projection("name", "biggestCity"),
                                                                  projection("pop", "biggestPop")),
                                                       projection("smallestCity",
                                                                  projection("name", "smallestCity"),
                                                                  projection("pop", "smallestPop")));

        Iterator<State> iterator = pipeline.aggregate(State.class);
        try {
            Map<String, State> states = new HashMap<String, State>();
            while (iterator.hasNext()) {
                State state = iterator.next();
{'fraction_non_alphanumeric': 0.11380195683852612, 'fraction_numerical': 0.005273749219346333, 'mean_word_length': 1.9003823706983296, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 18, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 98, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1568311', 'n_tokens_mistral': 3630, 'n_tokens_neox': 3432, 'n_words': 670}