test for request cookies and headers

master
yihua.huang 2017-04-08 23:13:16 +08:00
parent db67db8103
commit a1ae632b62
3 changed files with 86 additions and 28 deletions

View File

@ -1,14 +1,10 @@
package us.codecraft.webmagic; package us.codecraft.webmagic;
import org.apache.http.Header;
import org.apache.http.cookie.Cookie;
import us.codecraft.webmagic.model.HttpRequestBody; import us.codecraft.webmagic.model.HttpRequestBody;
import us.codecraft.webmagic.utils.Experimental; import us.codecraft.webmagic.utils.Experimental;
import java.io.Serializable; import java.io.Serializable;
import java.util.ArrayList;
import java.util.HashMap; import java.util.HashMap;
import java.util.List;
import java.util.Map; import java.util.Map;
/** /**
@ -38,9 +34,9 @@ public class Request implements Serializable {
/** /**
* cookies for current url, if not set use Site's cookies * cookies for current url, if not set use Site's cookies
*/ */
private List<Cookie> cookies=new ArrayList<Cookie>(); private Map<String, String> cookies = new HashMap<String, String>();
private List<Header> headers=new ArrayList<Header>(); private Map<String, String> headers = new HashMap<String, String>();
/** /**
* Priority of the request.<br> * Priority of the request.<br>
@ -137,11 +133,21 @@ public class Request implements Serializable {
return method != null ? method.equals(request.method) : request.method == null; return method != null ? method.equals(request.method) : request.method == null;
} }
public List<Cookie> getCookies() { public Request addCookie(String name, String value) {
cookies.put(name, value);
return this;
}
public Request addHeader(String name, String value) {
headers.put(name, value);
return this;
}
public Map<String, String> getCookies() {
return cookies; return cookies;
} }
public List<Header> getHeaders() { public Map<String, String> getHeaders() {
return headers; return headers;
} }

View File

@ -1,7 +1,5 @@
package us.codecraft.webmagic.downloader; package us.codecraft.webmagic.downloader;
import org.apache.commons.collections.CollectionUtils;
import org.apache.http.Header;
import org.apache.http.HttpHost; import org.apache.http.HttpHost;
import org.apache.http.auth.AuthState; import org.apache.http.auth.AuthState;
import org.apache.http.auth.UsernamePasswordCredentials; import org.apache.http.auth.UsernamePasswordCredentials;
@ -11,14 +9,15 @@ import org.apache.http.client.config.RequestConfig;
import org.apache.http.client.methods.HttpUriRequest; import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.client.methods.RequestBuilder; import org.apache.http.client.methods.RequestBuilder;
import org.apache.http.client.protocol.HttpClientContext; import org.apache.http.client.protocol.HttpClientContext;
import org.apache.http.cookie.Cookie;
import org.apache.http.entity.ByteArrayEntity; import org.apache.http.entity.ByteArrayEntity;
import org.apache.http.impl.auth.BasicScheme; import org.apache.http.impl.auth.BasicScheme;
import org.apache.http.impl.client.BasicCookieStore; import org.apache.http.impl.client.BasicCookieStore;
import org.apache.http.impl.cookie.BasicClientCookie;
import us.codecraft.webmagic.Request; import us.codecraft.webmagic.Request;
import us.codecraft.webmagic.Site; import us.codecraft.webmagic.Site;
import us.codecraft.webmagic.proxy.Proxy; import us.codecraft.webmagic.proxy.Proxy;
import us.codecraft.webmagic.utils.HttpConstant; import us.codecraft.webmagic.utils.HttpConstant;
import us.codecraft.webmagic.utils.UrlUtils;
import java.util.Map; import java.util.Map;
@ -45,10 +44,12 @@ public class HttpUriRequestConverter {
authState.update(new BasicScheme(), new UsernamePasswordCredentials(proxy.getUsername(), proxy.getPassword())); authState.update(new BasicScheme(), new UsernamePasswordCredentials(proxy.getUsername(), proxy.getPassword()));
httpContext.setAttribute(HttpClientContext.PROXY_AUTH_STATE, authState); httpContext.setAttribute(HttpClientContext.PROXY_AUTH_STATE, authState);
} }
if (request.getCookies() != null && CollectionUtils.isNotEmpty(request.getCookies())) { if (request.getCookies() != null && !request.getCookies().isEmpty()) {
CookieStore cookieStore = new BasicCookieStore(); CookieStore cookieStore = new BasicCookieStore();
for (Cookie c : request.getCookies()) { for (Map.Entry<String, String> cookieEntry : request.getCookies().entrySet()) {
cookieStore.addCookie(c); BasicClientCookie cookie1 = new BasicClientCookie(cookieEntry.getKey(), cookieEntry.getValue());
cookie1.setDomain(UrlUtils.removePort(UrlUtils.getDomain(request.getUrl())));
cookieStore.addCookie(cookie1);
} }
httpContext.setCookieStore(cookieStore); httpContext.setCookieStore(cookieStore);
} }
@ -76,9 +77,9 @@ public class HttpUriRequestConverter {
} }
requestBuilder.setConfig(requestConfigBuilder.build()); requestBuilder.setConfig(requestConfigBuilder.build());
HttpUriRequest httpUriRequest = requestBuilder.build(); HttpUriRequest httpUriRequest = requestBuilder.build();
if (request.getHeaders() != null && CollectionUtils.isNotEmpty(request.getHeaders())) { if (request.getHeaders() != null && !request.getHeaders().isEmpty()) {
for (Header h : request.getHeaders()) { for (Map.Entry<String, String> header : request.getHeaders().entrySet()) {
httpUriRequest.setHeader(h); httpUriRequest.addHeader(header.getKey(), header.getValue());
} }
} }
return httpUriRequest; return httpUriRequest;

View File

@ -34,7 +34,7 @@ import static org.junit.Assert.assertTrue;
*/ */
public class HttpClientDownloaderTest { public class HttpClientDownloaderTest {
public static final String PAGE_ALWAYS_NOT_EXISTS = "http://localhost:13421/404"; public static final String PAGE_ALWAYS_NOT_EXISTS = "http://localhost:13423/404";
@Test @Test
public void testDownloader() { public void testDownloader() {
@ -63,7 +63,7 @@ public class HttpClientDownloaderTest {
@Test @Test
public void testGetHtmlCharset() throws Exception { public void testGetHtmlCharset() throws Exception {
HttpServer server = httpserver(12306); HttpServer server = httpserver(13423);
server.get(by(uri("/header"))).response(header("Content-Type", "text/html; charset=gbk")); server.get(by(uri("/header"))).response(header("Content-Type", "text/html; charset=gbk"));
server.get(by(uri("/meta4"))).response(with(text("<html>\n" + server.get(by(uri("/meta4"))).response(with(text("<html>\n" +
" <head>\n" + " <head>\n" +
@ -80,11 +80,11 @@ public class HttpClientDownloaderTest {
Runner.running(server, new Runnable() { Runner.running(server, new Runnable() {
@Override @Override
public void run() { public void run() {
String charset = getCharsetByUrl("http://127.0.0.1:12306/header"); String charset = getCharsetByUrl("http://127.0.0.1:13423/header");
assertEquals(charset, "gbk"); assertEquals(charset, "gbk");
charset = getCharsetByUrl("http://127.0.0.1:12306/meta4"); charset = getCharsetByUrl("http://127.0.0.1:13423/meta4");
assertEquals(charset, "gbk"); assertEquals(charset, "gbk");
charset = getCharsetByUrl("http://127.0.0.1:12306/meta5"); charset = getCharsetByUrl("http://127.0.0.1:13423/meta5");
assertEquals(charset, "gbk"); assertEquals(charset, "gbk");
} }
@ -114,7 +114,7 @@ public class HttpClientDownloaderTest {
@Test @Test
public void test_selectRequestMethod() throws Exception { public void test_selectRequestMethod() throws Exception {
HttpServer server = httpserver(12306); HttpServer server = httpserver(13423);
server.get(eq(query("q"), "webmagic")).response("get"); server.get(eq(query("q"), "webmagic")).response("get");
server.post(eq(form("q"), "webmagic")).response("post"); server.post(eq(form("q"), "webmagic")).response("post");
server.put(eq(form("q"), "webmagic")).response("put"); server.put(eq(form("q"), "webmagic")).response("put");
@ -127,7 +127,7 @@ public class HttpClientDownloaderTest {
@Override @Override
public void run() throws Exception { public void run() throws Exception {
Request request = new Request(); Request request = new Request();
request.setUrl("http://127.0.0.1:12306/search?q=webmagic"); request.setUrl("http://127.0.0.1:13423/search?q=webmagic");
request.setMethod(HttpConstant.Method.GET); request.setMethod(HttpConstant.Method.GET);
Map<String,Object> params = new HashedMap(); Map<String,Object> params = new HashedMap();
params.put("q","webmagic"); params.put("q","webmagic");
@ -142,7 +142,7 @@ public class HttpClientDownloaderTest {
request.setMethod(HttpConstant.Method.TRACE); request.setMethod(HttpConstant.Method.TRACE);
httpUriRequest = httpUriRequestConverter.convert(request, site, null).getHttpUriRequest(); httpUriRequest = httpUriRequestConverter.convert(request, site, null).getHttpUriRequest();
assertThat(EntityUtils.toString(HttpClients.custom().build().execute(httpUriRequest).getEntity())).isEqualTo("trace"); assertThat(EntityUtils.toString(HttpClients.custom().build().execute(httpUriRequest).getEntity())).isEqualTo("trace");
request.setUrl("http://127.0.0.1:12306/search"); request.setUrl("http://127.0.0.1:13423/search");
request.setMethod(HttpConstant.Method.POST); request.setMethod(HttpConstant.Method.POST);
request.setRequestBody(HttpRequestBody.form(params, "utf-8")); request.setRequestBody(HttpRequestBody.form(params, "utf-8"));
httpUriRequest = httpUriRequestConverter.convert(request, site, null).getHttpUriRequest(); httpUriRequest = httpUriRequestConverter.convert(request, site, null).getHttpUriRequest();
@ -154,16 +154,67 @@ public class HttpClientDownloaderTest {
}); });
} }
@Test
public void test_set_request_cookie() throws Exception {
HttpServer server = httpserver(13423);
server.get(eq(cookie("cookie"), "cookie-webmagic")).response("ok");
Runner.running(server, new Runnable() {
@Override
public void run() throws Exception {
HttpClientDownloader httpClientDownloader = new HttpClientDownloader();
Request request = new Request();
request.setUrl("http://127.0.0.1:13423");
request.addCookie("cookie","cookie-webmagic");
Page page = httpClientDownloader.download(request, Site.me().toTask());
assertThat(page.getRawText()).isEqualTo("ok");
}
});
}
@Test
public void test_set_request_header() throws Exception {
HttpServer server = httpserver(13423);
server.get(eq(header("header"), "header-webmagic")).response("ok");
Runner.running(server, new Runnable() {
@Override
public void run() throws Exception {
HttpClientDownloader httpClientDownloader = new HttpClientDownloader();
Request request = new Request();
request.setUrl("http://127.0.0.1:13423");
request.addHeader("header","header-webmagic");
Page page = httpClientDownloader.download(request, Site.me().toTask());
assertThat(page.getRawText()).isEqualTo("ok");
}
});
}
@Test
public void test_set_site_cookie() throws Exception {
HttpServer server = httpserver(13423);
server.get(eq(cookie("cookie"), "cookie-webmagic")).response("ok");
Runner.running(server, new Runnable() {
@Override
public void run() throws Exception {
HttpClientDownloader httpClientDownloader = new HttpClientDownloader();
Request request = new Request();
request.setUrl("http://127.0.0.1:13423");
Site site = Site.me().addCookie("cookie", "cookie-webmagic").setDomain("127.0.0.1");
Page page = httpClientDownloader.download(request, site.toTask());
assertThat(page.getRawText()).isEqualTo("ok");
}
});
}
@Test @Test
public void test_download_when_task_is_null() throws Exception { public void test_download_when_task_is_null() throws Exception {
HttpServer server = httpserver(12306); HttpServer server = httpserver(13423);
server.response("foo"); server.response("foo");
Runner.running(server, new Runnable() { Runner.running(server, new Runnable() {
@Override @Override
public void run() throws Exception { public void run() throws Exception {
final HttpClientDownloader httpClientDownloader = new HttpClientDownloader(); final HttpClientDownloader httpClientDownloader = new HttpClientDownloader();
Request request = new Request(); Request request = new Request();
request.setUrl("http://127.0.0.1:12306/"); request.setUrl("http://127.0.0.1:13423/");
Page page = httpClientDownloader.download(request, Site.me().toTask()); Page page = httpClientDownloader.download(request, Site.me().toTask());
assertThat(page.getRawText()).isEqualTo("foo"); assertThat(page.getRawText()).isEqualTo("foo");
} }