Examples with Attribute datawave.query.attributes.Attribute used on opensource projects

Example 21 with Attribute

use of datawave.query.attributes.Attribute in project datawave by NationalSecurityAgency.

the class HitsAreAlwaysIncludedTest method runTestQuery.

protected void runTestQuery(Connector connector, String queryString, Date startDate, Date endDate, Map<String, String> extraParms, Collection<String> expectedHits, Collection<String> goodResults) throws Exception {
    QueryImpl settings = new QueryImpl();
    settings.setBeginDate(startDate);
    settings.setEndDate(endDate);
    settings.setPagesize(Integer.MAX_VALUE);
    settings.setQueryAuthorizations(auths.serialize());
    settings.setQuery(queryString);
    settings.setParameters(extraParms);
    settings.setId(UUID.randomUUID());
    log.debug("query: " + settings.getQuery());
    log.debug("logic: " + settings.getQueryLogicName());
    GenericQueryConfiguration config = logic.initialize(connector, settings, authSet);
    logic.setupQuery(config);
    Set<Document> docs = new HashSet<>();
    for (Entry<Key, Value> entry : logic) {
        Document d = deserializer.apply(entry).getValue();
        log.trace(entry.getKey() + " => " + d);
        docs.add(d);
        Attribute hitAttribute = d.get(JexlEvaluation.HIT_TERM_FIELD);
        if (hitAttribute instanceof Attributes) {
            Attributes attributes = (Attributes) hitAttribute;
            for (Attribute attr : attributes.getAttributes()) {
                if (attr instanceof Content) {
                    Content content = (Content) attr;
                    Assert.assertTrue(expectedHits.remove(content.getContent()));
                }
            }
        } else if (hitAttribute instanceof Content) {
            Content content = (Content) hitAttribute;
            Assert.assertTrue(content.getContent() + " is not an expected hit", expectedHits.remove(content.getContent()));
        } else {
            Assert.fail("Did not find hit term field");
        }
        Assert.assertTrue(expectedHits + " expected hits was not empty", expectedHits.isEmpty());
        // remove from goodResults as we find the expected return fields
        log.debug("goodResults: " + goodResults);
        Map<String, Attribute<? extends Comparable<?>>> dictionary = d.getDictionary();
        log.debug("dictionary:" + dictionary);
        for (Entry<String, Attribute<? extends Comparable<?>>> dictionaryEntry : dictionary.entrySet()) {
            Attribute<? extends Comparable<?>> attribute = dictionaryEntry.getValue();
            if (attribute instanceof Attributes) {
                for (Attribute attr : ((Attributes) attribute).getAttributes()) {
                    String toFind = dictionaryEntry.getKey() + ":" + attr;
                    boolean found = goodResults.remove(toFind);
                    if (found)
                        log.debug("removed " + toFind);
                    else
                        log.debug("Did not remove " + toFind);
                }
            } else {
                String toFind = dictionaryEntry.getKey() + ":" + dictionaryEntry.getValue();
                boolean found = goodResults.remove(toFind);
                if (found)
                    log.debug("removed " + toFind);
                else
                    log.debug("Did not remove " + toFind);
            }
        }
        Assert.assertTrue(goodResults + " good results was not empty", goodResults.isEmpty());
    }
    Assert.assertTrue("No docs were returned!", !docs.isEmpty());
}