<commit_before>/**
 * 
 */
package edu.jhu.hlt.rebar.accumulo;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.fail;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import org.apache.accumulo.core.client.BatchScanner;
import org.apache.accumulo.core.client.IteratorSetting;
import org.apache.accumulo.core.client.Scanner;
import org.apache.accumulo.core.client.ScannerBase;
import org.apache.accumulo.core.client.TableNotFoundException;
import org.apache.accumulo.core.data.Key;
import org.apache.accumulo.core.data.Range;
import org.apache.accumulo.core.data.Value;
import org.apache.accumulo.core.iterators.user.WholeRowIterator;
import org.apache.hadoop.io.Text;
import org.apache.thrift.TException;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import edu.jhu.hlt.concrete.Communication;
import edu.jhu.hlt.concrete.LangId;
import edu.jhu.hlt.concrete.LanguagePrediction;
import edu.jhu.hlt.concrete.Stage;
import edu.jhu.hlt.concrete.StageType;
import edu.jhu.hlt.rebar.Constants;
import edu.jhu.hlt.rebar.Configuration;
import edu.jhu.hlt.rebar.RebarException;
import edu.jhu.hlt.rebar.Util;

/**
 * @author max
 * 
 */
public class TestRebarReader extends AbstractAccumuloTest {

  private RebarReader rr;

  /**
   * @throws java.lang.Exception
   */
  @Before
  public void setUp() throws Exception {
    this.initialize();
    this.rr = new RebarReader(this.conn);
  }

  /**
   * @throws java.lang.Exception
   */
  @After
  public void tearDown() throws Exception {
    this.rr.close();
  }

  /**
   * Test method for {@link edu.jhu.hlt.rebar.accumulo.RebarReader#getAnnotatedDocuments(com.maxjthomas.dumpster.Stage)}.
   * 
   * @throws Exception
   * @throws RebarException
   */
  @Test
  public void testGetAnnotatedDocuments() throws RebarException, Exception {
    int nDocs = 3;
    List<Communication> docList = this.ingestDocuments(nDocs);

    Set<Communication> docsWithLid = new HashSet<>();

    List<LangId> langIdList = new ArrayList<>();
    Stage s = generateTestStage();
    try (RebarAnnotator ra = new RebarAnnotator(this.conn);) {
      for (Communication d : docList) {
        LangId mockLid = generateLangId(d);
        ra.addLanguageId(d, s, mockLid);
        langIdList.add(mockLid);
        Communication newDoc = new Communication(d);
        newDoc.lid = mockLid;
        docsWithLid.add(newDoc);
      }
    }

    int annotatedDocs = 0;
    Set<String> idSet;
    try (RebarStageHandler ash = new RebarStageHandler(this.conn);) {
      annotatedDocs = ash.getAnnotatedDocumentCount(s);
      idSet = ash.getAnnotatedDocumentIds(s);
    }

    assertEquals("Should get n annotated docs: (n = " + nDocs + ")", nDocs, annotatedDocs);
    BatchScanner bsc = this.rr.createScanner(s, idSet);
    assertEquals("Should get " + nDocs + " entries in this batch scanner.", 3, Util.countIteratorResults(bsc.iterator()));
    bsc.close();

    List<Communication> fetchedDocs = this.rr.getAnnotatedCommunications(s);
    assertEquals("Documents with LID should be the same.", new ArrayList<>(docsWithLid), fetchedDocs);
  }

  /**
   * Test method for {@link edu.jhu.hlt.rebar.accumulo.RebarReader#getAnnotatedDocuments(com.maxjthomas.dumpster.Stage)}.
   * 
   * @throws Exception
   * @throws RebarException
   */
  @Test
  public void testGetAnnotatedDocumentsStageDependency() throws RebarException, Exception {
    // create the stages
    Stage stageA = new Stage("stage_max_lid_test", "Testing stage for LID", Util.getCurrentUnixTime(), new HashSet<String>(), StageType.LANG_ID);
    Set<String> stageBDeps = new HashSet<>();
    stageBDeps.add(stageA.name);
    Stage stageB = new Stage("stage_max_lp_test", "Testing stage for LP", Util.getCurrentUnixTime(), stageBDeps, StageType.LANG_PRED);

    // ingest documents
    int nDocs = 3;
    List<Communication> docList = this.ingestDocuments(nDocs);

    // annotate documents
    try (RebarAnnotator ra = new RebarAnnotator(this.conn);) {
      for (Communication d : docList) {
        LangId mockLid = generateLangId(d);
        ra.addLanguageId(d, stageA, mockLid);
      }
    }

    List<Communication> fetchedDocs = this.rr.getAnnotatedCommunications(stageA);

    Set<Communication> docsWithLid = new HashSet<>();
    try (RebarAnnotator ra = new RebarAnnotator(this.conn);) {
      // add language prediction annotation
      for (Communication d : fetchedDocs) {
        LangId lid = d.lid;
        LanguagePrediction lp = new LanguagePrediction();
        Entry<String, Double> e = lid.languageToProbabilityMap.entrySet().iterator().next();

        switch (e.getKey()) {
        case "eng":
          lp.predictedLanguage = "English";
          break;
        case "fra":
          lp.predictedLanguage = "French";
          break;
        case "spa":
          lp.predictedLanguage = "Spanish";
          break;
        default:
          lp.predictedLanguage = "Unknown";
        }

        ra.addLanguagePrediction(d, stageB, lp);
        Communication newDoc = new Communication(d);
        newDoc.language = lp;
        docsWithLid.add(newDoc);
      }
    }
    
    int annotatedDocs = 0;
    try (RebarStageHandler ash = new RebarStageHandler(this.conn);) {
      annotatedDocs = ash.getAnnotatedDocumentCount(stageB);
    }

    assertEquals("Should get n annotated docs: (n = " + nDocs + ")", nDocs, annotatedDocs);

    fetchedDocs = this.rr.getAnnotatedCommunications(stageB);
    assertEquals("Documents with LID should be the same.", new ArrayList<>(docsWithLid), fetchedDocs);
  }
  
  /**
   * Test method for {@link edu.jhu.hlt.rebar.accumulo.RebarReader#getAnnotatedDocuments(com.maxjthomas.dumpster.Stage)}.
   * 
   * @throws Exception
   * @throws RebarException
   */
  @Test
  public void testGetAnnotatedDocumentsManyStages() throws RebarException, Exception {
    // create the stages<commit_msg>fix a non-deterministic test<commit_after>/**
 * 
 */
package edu.jhu.hlt.rebar.accumulo;

import static org.junit.Assert.*;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import org.apache.accumulo.core.client.BatchScanner;
import org.apache.accumulo.core.client.IteratorSetting;
import org.apache.accumulo.core.client.Scanner;
import org.apache.accumulo.core.client.ScannerBase;
import org.apache.accumulo.core.client.TableNotFoundException;
import org.apache.accumulo.core.data.Key;
import org.apache.accumulo.core.data.Range;
import org.apache.accumulo.core.data.Value;
import org.apache.accumulo.core.iterators.user.WholeRowIterator;
import org.apache.hadoop.io.Text;
import org.apache.thrift.TException;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import edu.jhu.hlt.concrete.Communication;
import edu.jhu.hlt.concrete.LangId;
import edu.jhu.hlt.concrete.LanguagePrediction;
import edu.jhu.hlt.concrete.Stage;
import edu.jhu.hlt.concrete.StageType;
import edu.jhu.hlt.rebar.Constants;
import edu.jhu.hlt.rebar.Configuration;
import edu.jhu.hlt.rebar.RebarException;
import edu.jhu.hlt.rebar.Util;

/**
 * @author max
 * 
 */
public class TestRebarReader extends AbstractAccumuloTest {

  private RebarReader rr;

  /**
   * @throws java.lang.Exception
   */
  @Before
  public void setUp() throws Exception {
    this.initialize();
    this.rr = new RebarReader(this.conn);
  }

  /**
   * @throws java.lang.Exception
   */
  @After
  public void tearDown() throws Exception {
    this.rr.close();
  }

  /**
   * Test method for {@link edu.jhu.hlt.rebar.accumulo.RebarReader#getAnnotatedDocuments(com.maxjthomas.dumpster.Stage)}.
   * 
   * @throws Exception
   * @throws RebarException
   */
  @Test
  public void testGetAnnotatedDocuments() throws RebarException, Exception {
    int nDocs = 3;
    List<Communication> docList = this.ingestDocuments(nDocs);

    Set<Communication> docsWithLid = new HashSet<>();

    List<LangId> langIdList = new ArrayList<>();
    Stage s = generateTestStage();
    try (RebarAnnotator ra = new RebarAnnotator(this.conn);) {
      for (Communication d : docList) {
        LangId mockLid = generateLangId(d);
        ra.addLanguageId(d, s, mockLid);
        langIdList.add(mockLid);
        Communication newDoc = new Communication(d);
        newDoc.lid = mockLid;
        docsWithLid.add(newDoc);
      }
    }

    int annotatedDocs = 0;
    Set<String> idSet;
    try (RebarStageHandler ash = new RebarStageHandler(this.conn);) {
      annotatedDocs = ash.getAnnotatedDocumentCount(s);
      idSet = ash.getAnnotatedDocumentIds(s);
    }

    assertEquals("Should get n annotated docs: (n = " + nDocs + ")", nDocs, annotatedDocs);
    BatchScanner bsc = this.rr.createScanner(s, idSet);
    assertEquals("Should get " + nDocs + " entries in this batch scanner.", 3, Util.countIteratorResults(bsc.iterator()));
    bsc.close();

    List<Communication> fetchedDocs = this.rr.getAnnotatedCommunications(s);
    assertEquals("Documents with LID should be the same.", new ArrayList<>(docsWithLid), fetchedDocs);
  }

  /**
   * Test method for {@link edu.jhu.hlt.rebar.accumulo.RebarReader#getAnnotatedDocuments(com.maxjthomas.dumpster.Stage)}.
   * 
   * @throws Exception
   * @throws RebarException
   */
  @Test
  public void testGetAnnotatedDocumentsStageDependency() throws RebarException, Exception {
    // create the stages
    Stage stageA = new Stage("stage_max_lid_test", "Testing stage for LID", Util.getCurrentUnixTime(), new HashSet<String>(), StageType.LANG_ID);
    Set<String> stageBDeps = new HashSet<>();
    stageBDeps.add(stageA.name);
    Stage stageB = new Stage("stage_max_lp_test", "Testing stage for LP", Util.getCurrentUnixTime(), stageBDeps, StageType.LANG_PRED);

    // ingest documents
    int nDocs = 3;
    List<Communication> docList = this.ingestDocuments(nDocs);

    // annotate documents
    try (RebarAnnotator ra = new RebarAnnotator(this.conn);) {
      for (Communication d : docList) {
        LangId mockLid = generateLangId(d);
        ra.addLanguageId(d, stageA, mockLid);
      }
    }

    List<Communication> fetchedDocs = this.rr.getAnnotatedCommunications(stageA);

    Set<Communication> docsWithLid = new HashSet<>();
    try (RebarAnnotator ra = new RebarAnnotator(this.conn);) {
      // add language prediction annotation
      for (Communication d : fetchedDocs) {
        LangId lid = d.lid;
        LanguagePrediction lp = new LanguagePrediction();
        Entry<String, Double> e = lid.languageToProbabilityMap.entrySet().iterator().next();

        switch (e.getKey()) {
        case "eng":
          lp.predictedLanguage = "English";
          break;
        case "fra":
          lp.predictedLanguage = "French";
          break;
        case "spa":
          lp.predictedLanguage = "Spanish";
          break;
        default:
          lp.predictedLanguage = "Unknown";
        }

        ra.addLanguagePrediction(d, stageB, lp);
        Communication newDoc = new Communication(d);
        newDoc.language = lp;
        docsWithLid.add(newDoc);
      }
    }
    
    int annotatedDocs = 0;
    try (RebarStageHandler ash = new RebarStageHandler(this.conn);) {
      annotatedDocs = ash.getAnnotatedDocumentCount(stageB);
    }

    assertEquals("Should get n annotated docs: (n = " + nDocs + ")", nDocs, annotatedDocs);

    fetchedDocs = this.rr.getAnnotatedCommunications(stageB);
    assertTrue("Documents with LID should be the same.", fetchedDocs.containsAll(docsWithLid));
  }
  
  /**
   * Test method for {@link edu.jhu.hlt.rebar.accumulo.RebarReader#getAnnotatedDocuments(com.maxjthomas.dumpster.Stage)}.
   * 
   * @throws Exception
   * @throws RebarException
   */
  @Test
  public void testGetAnnotatedDocumentsManyStages() throws RebarException, Exception {
    // create the stages
{'fraction_non_alphanumeric': 0.11914611947708093, 'fraction_numerical': 0.0008274036074797286, 'mean_word_length': 3.4866369710467704, 'pattern_counts': {'":': 6, '<': 36, '<?xml version=': 0, '>': 36, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4891547', 'n_tokens_mistral': 4080, 'n_tokens_neox': 3910, 'n_words': 785}