DaoProvider.java

package me.xhh.hector06;

import me.xhh.hector06.DBA;
import me.xhh.hector06.DBConfig;
import me.xhh.hector06.DBA.ColumnFamily;

import org.testng.annotations.AfterMethod;

/**
 * Provide database configurations for testing environment.
 */
public abstract class DaoProvider {

	protected static final DBConfig DB_CONFIG = new DBConfig("localhost", 9160, "demo-test");
	private DBA dba;

	@SuppressWarnings("unchecked")
	protected <T> T prepareDao(Class<?> cls) throws Exception {
		try {
			dba = (DBA) cls.newInstance();
			dba.setDbConfig(DB_CONFIG);
			return (T) dba;
		} catch (Exception e) {
			throw e;
		}
	}

	private void removeAllKeys(ColumnFamily cf) {
		for (String key : dba.getAllKeys(cf))
			dba.deleteKey(cf, key);
	}

	@AfterMethod
	protected void clearData() {
		for (ColumnFamily cf : getColumnFamilies())
			removeAllKeys(cf);
	}

	/**
	 * @return the (super) column families whose data should be cleared after each test
	 */
	protected abstract ColumnFamily[] getColumnFamilies();

}