import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.*;
import java.util.regex.Matcher; public class SQLProvider implements Iterator<Object[]> {
List<Map<String, String>> sqlList = new ArrayList();
int rowNum=0; //行数
int curRowNo=0; //当前行数
String sql = null;
String category = ""; /**
* 在TestNG中由@DataProvider(dataProvider = "name")修饰的方法
* 取SQL数据时,调用此类构造方法(此方法会得到列名并将当前行移到下以后)执行后,转发到
* TestNG自己的方法中去,然后由它们调用此类实现的hasNext()、next()方法
* 得到一行数据,然后返回给由@Test(dataProvider = "name")修饰的方法,如此
* 反复到数据读完为止
* @param tablename
*/
public SQLProvider(String tablename, String ... key) {
String ip = null;
String user = null;
String pw = null;
String db = null;
Properties prop = new Properties();
try {
File directory=new File(".");
String path=".src.test.resources.properties.";
String absolutePath=directory.getCanonicalPath()+path.replaceAll("\\.", Matcher.quoteReplacement("\\"))+"conf.properties";
absolutePath = absolutePath.replace("\\",File.separator).replace("/",File.separator);
InputStream in = new FileInputStream(absolutePath);
prop.load(in);
ip = prop.getProperty("SQLProviderIp").trim();
user = prop.getProperty("SQLProviderUser").trim();
pw = prop.getProperty("SQLProviderPw").trim();
db = prop.getProperty("SQLProviderDB").trim();
} catch (IOException e) {
e.printStackTrace();
}
if (key.length > 0){
for (int i = 0 ;i < key.length;i++){
category += "'"+key[i]+"',";
}
category = category.substring(0,category.length()-1);
sql = "select * from "+tablename+" where category in ("+category+")";
}else{
sql = "select * from " + tablename;
}
sqlList = getConnection(sql,ip,user,pw,db);
this.rowNum = sqlList.size();
} @Override
public boolean hasNext() {
// TODO Auto-generated method stub
if(rowNum==0||curRowNo>=rowNum){
return false;
}else{
return true;
}
}
/**
* 获取一组参数,即一行数据
*/
@Override
public Object[] next() {
// TODO Auto-generated method stub
Map<String,String> s=new TreeMap<String,String>();
s=sqlList.get(curRowNo);
Object[] d=new Object[1];
d[0]=s;
this.curRowNo++;
return d;
} @Override
public void remove() {
// TODO Auto-generated method stub
throw new UnsupportedOperationException("remove unsupported");
} public static List<Map<String,String>> getConnection(String sql,String ip,String user,String pw,String db) {
Connection conn = null;
List<Map<String, String>> list = new ArrayList();
try {
Class.forName("com.mysql.jdbc.Driver");
String url = "jdbc:mysql://" + ip + "/"+db+"?user=" + user
+ "&password=" + pw
+ "&useUnicode=true&characterEncoding=UTF8";
conn = DriverManager.getConnection(url);
Statement stmt = conn.createStatement();
ResultSet result = stmt.executeQuery(sql);
while (result.next()) {
Map<String, String> m = new HashMap<String, String>();
for (int i = 2; i <= result.getMetaData().getColumnCount(); i++) {
String k = result.getMetaData().getColumnName(i);
String v = result.getString(i);
m.put(k,v);
}
list.add(m);
}
stmt.close();
conn.close();
} catch (Exception e) {
e.printStackTrace();
} return list; } /* public static void main (String []args) throws IOException {
ConnDB d = new ConnDB();
d.setDb("dataprovider");
d.setDbUrl("localhost");
d.setPwd("shen1986");
d.setUserName("root");
d.ConnDB();
System.out.println(getConnection("select * from zfk",d.getConn()));
}*/ }
数据库配置文件:
#DataProvider SQL Config
#----------------------------------------------------------------------------#
SQLProviderIp=10.253.43.98
SQLProviderUser=root
SQLProviderPw=qwER!@#$1234
SQLProviderDB=dataprovider
testng中使用:
@DataProvider(name="periodnum")
public Iterator<Object[]> periodnum() throws IOException {
return (Iterator<Object[]>)new SQLProvider("mmd","proidnum3","proidnum6","proidnum12");
}