DistributedTest.java

package test.distributed;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Random;

import org.testng.Assert;
import org.testng.ITestResult;
import org.testng.TestListenerAdapter;
import org.testng.TestNG;
import org.testng.annotations.Configuration;
import org.testng.annotations.Test;
import org.testng.remote.SuiteDispatcher;
import org.testng.xml.XmlClass;
import org.testng.xml.XmlSuite;
import org.testng.xml.XmlTest;

import test.BaseDistributedTest;

import testhelper.OutputDirectoryPatch;

public class DistributedTest extends BaseDistributedTest {
  
  private List<Thread> m_hostThreads = new ArrayList<Thread>();

  protected Thread startSlave(final String filename) {
    Thread result = new Thread(new Runnable() {
      public void run() {
        TestNG.main(new String[] { "-slave", filename, "-d", OutputDirectoryPatch.getOutputDirectory() });       
      }
    });
    result.setName("Slave-" + filename);
    result.start();
    m_hostThreads.add(result);
    return result;
  }
  
  private File createMasterProperties(String strategy)
  throws IOException 
  {
    String fileName = "remote";
    
    File result = File.createTempFile(fileName, ".properties");
    result.deleteOnExit();
    Properties p = new Properties();
    p.setProperty("testng.hosts", "localhost:" + m_ports[0] + " localhost:" + m_ports[1]);
    p.setProperty("testng.master.strategy", strategy);
    p.setProperty("testng.verbose", "0");
    p.setProperty("testng.master.adpter", "org.testng.remote.adapter.DefaultMastertAdapter");
    FileOutputStream fos = new FileOutputStream(result);
    p.store(fos, "Automatically generated by tests");
    fos.close();    
    
    return result;
  }
  
  private File createSlaveProperties(String port)
  throws IOException 
  {
    String fileName = "remote";
    
    File result = File.createTempFile(fileName, ".properties");
    result.deleteOnExit();
    Properties p = new Properties();
    p.setProperty("testng.verbose", "0");
    p.setProperty("slave.port", port);
    p.setProperty("testng.slave.adpter", "org.testng.remote.adapter.DefaultWorkerAdapter");
    FileOutputStream fos = new FileOutputStream(result);
    p.store(fos, "Automatically generated by tests");
    fos.close();    
    
    return result;
  }
  
  private XmlSuite createSuite(String name, Class[] classes) {
    XmlSuite result = new XmlSuite();
    result.setName(name);
    
    for (Class c : classes) {
      XmlTest test1 = new XmlTest(result);
      test1.setName(c.getName());
      XmlClass class1 = new XmlClass(c);
      test1.getXmlClasses().add(class1);
    }
    
    return result;
  }
  
//  @ Configuration(beforeTestClass = true) 
  private void startSlaves() throws IOException{
	  int port = new Random().nextInt(50000) + 2000;
	  m_ports = new String[] { Integer.toString(port), Integer.toString(port+1)};

	  File slaveFile = createSlaveProperties(m_ports[0]);
	  startSlave( slaveFile.getCanonicalPath());
	  
	  slaveFile = createSlaveProperties(m_ports[1]);
	  startSlave( slaveFile.getCanonicalPath());
  }

  private String[] m_ports = new String[2];

  public TestListenerAdapter twoHosts(String strategy) throws IOException {  
    TestNG tng = new TestNG();
    tng.setOutputDirectory(OutputDirectoryPatch.getOutputDirectory());
    
    File masterFile = createMasterProperties(strategy);
    tng.setMaster(masterFile.getAbsolutePath());
    
    XmlSuite suite = createSuite("DistributedSuite1", new Class[] { Test1.class, Test2.class });
    tng.setXmlSuites(Arrays.asList(new XmlSuite[] { suite }));
    
    TestListenerAdapter result = new TestListenerAdapter();
    tng.addListener(result);
    tng.run();
    
    String[] passed = {
        "f1", "f2"
    };
    String[] failed = {};
    String[] skipped = {};

    verifyTests("Passed", passed, toMap(result.getPassedTests()));
    verifyTests("Failed", failed, toMap(result.getFailedTests()));
    verifyTests("Skipped", skipped, toMap(result.getSkippedTests()));
    
    return result;
  }
  
  @Test
  public void twoHostsWithTestStrategy() throws IOException {
    startSlaves();
    TestListenerAdapter listener = twoHosts(SuiteDispatcher.STRATEGY_TEST);
    
    boolean found1 = false;
    boolean found2 = false;
    for (ITestResult tr : listener.getPassedTests()) {
      String host = tr.getHost();
      if (! found1) found1 = host.contains(m_ports[0]);
      if (! found2) found2 = host.contains(m_ports[1]);
    }
    Assert.assertTrue(found1, "No tests ran on port " + m_ports[0]);
    Assert.assertTrue(found2, "No tests ran on port " + m_ports[1]);    
  }
  
  @Test
  public void twoHostsWithSuiteStrategy() throws IOException {
    startSlaves();
    TestListenerAdapter listener = twoHosts( SuiteDispatcher.STRATEGY_SUITE);
  }

  private Map<String, ITestResult> toMap(List<ITestResult> results) {
    Map<String, ITestResult> result = new HashMap<String, ITestResult>();
    for (ITestResult tr : results) {
      result.put(tr.getName(), tr);
    }
    
    return result;
  }
  private void ppp(String string) {
    System.out.println("[DistributedTest] " + string);
  }
}