PerformanceMessenger.java

package jetbrains.mps.testbench;

/*Generated by MPS */

import org.junit.rules.TestRule;
import java.util.Map;
import java.util.LinkedHashMap;
import org.junit.runners.model.Statement;
import org.junit.runner.Description;
import org.junit.Assert;
import java.io.File;
import org.jdom.Element;
import org.jdom.Document;
import jetbrains.mps.util.JDOMUtil;
import java.io.IOException;
import org.jdom.JDOMException;

/**
 * Integrates with TeamCity's statistics. Contribute statistical values with teamcity-info.xml
 * To use, add a field annotated with {@link org.junit.Rule } or {@link org.junit.ClassRule } to your test
 * Evgeny Gryaznov, Oct 7, 2010
 */
public class PerformanceMessenger implements TestRule {
  private Map<String, Long> mySingleValues = new LinkedHashMap<String, Long>();
  private Map<String, long[]> myPercentValues = new LinkedHashMap<String, long[]>();
  private final String myKeyPrefix;

  public PerformanceMessenger() {
    myKeyPrefix = null;
  }

  public PerformanceMessenger(String keyPrefix) {
    myKeyPrefix = keyPrefix;
  }

  public Statement apply(final Statement base, Description description) {
    return new Statement() {
      public void evaluate() throws Throwable {
        base.evaluate();
        generateReport();
      }
    };
  }

  private String fullKey(String key) {
    return (myKeyPrefix == null ? key : myKeyPrefix + key);
  }

  public synchronized void report(String key, long value) {
    String fullKey = fullKey(key);
    Assert.assertFalse(myPercentValues.containsKey(fullKey));
    Long l = mySingleValues.remove(fullKey);
    if (l != null) {
      value += l;
    }
    mySingleValues.put(fullKey, value);
  }

  public synchronized void reportPercent(String key, long amount, long total) {
    String fullKey = fullKey(key);
    Assert.assertFalse(mySingleValues.containsKey(fullKey));
    long[] l = myPercentValues.get(fullKey);
    if (l != null) {
      l[0] += amount;
      l[1] += total;
    } else {
      myPercentValues.put(fullKey, new long[]{amount, total});
    }
  }

  public synchronized void generateReport() {
    try {
      if (mySingleValues.isEmpty() && myPercentValues.isEmpty()) {
        return;
      }
      File file = new File(System.getProperty("user.dir") + "/teamcity-info.xml");
      Element build;
      Document document;
      if (file.exists()) {
        document = JDOMUtil.loadDocument(file);
        build = document.getRootElement();
      } else {
        build = new Element("build");
        document = new Document(build);
      }

      for (Map.Entry<String, Long> e : mySingleValues.entrySet()) {
        Element child = new Element("statisticValue");
        child.setAttribute("key", e.getKey());
        child.setAttribute("value", Long.toString(e.getValue()));
        build.addContent(child);
      }
      for (Map.Entry<String, long[]> e : myPercentValues.entrySet()) {
        long amount = e.getValue()[0];
        long total = e.getValue()[1];
        Element child = new Element("statisticValue");
        child.setAttribute("key", e.getKey());
        child.setAttribute("value", Long.toString(amount * 100 / total));
        build.addContent(child);
      }
      JDOMUtil.writeDocument(document, file);
    } catch (IOException ex) {
      Assert.fail(ex.getMessage());
    } catch (JDOMException ex) {
      Assert.fail(ex.getMessage());
    }
  }
}