diff --git a/webmagic-avalon/webmagic-worker/pom.xml b/webmagic-avalon/webmagic-worker/pom.xml index f7833a0c..84e1d734 100644 --- a/webmagic-avalon/webmagic-worker/pom.xml +++ b/webmagic-avalon/webmagic-worker/pom.xml @@ -16,6 +16,18 @@ webmagic-avalon-common ${project.version} + + junit + junit + + + org.mockito + mockito-all + + + org.aspectj + aspectjrt + diff --git a/webmagic-avalon/webmagic-worker/src/main/java/us/codecraft/webmagic/worker/Worker.java b/webmagic-avalon/webmagic-worker/src/main/java/us/codecraft/webmagic/worker/Worker.java new file mode 100644 index 00000000..91867cc9 --- /dev/null +++ b/webmagic-avalon/webmagic-worker/src/main/java/us/codecraft/webmagic/worker/Worker.java @@ -0,0 +1,46 @@ +package us.codecraft.webmagic.worker; + +import us.codecraft.webmagic.Spider; +import us.codecraft.webmagic.utils.ThreadUtils; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutorService; + +/** + * @author code4crafter@gmail.com + */ +public class Worker { + + public static final int DEFAULT_POOL_SIZE = 10; + + private int poolSize; + + private ExecutorService executorService; + + private Map spiderMap; + + public Worker(int poolSize) { + this.poolSize = poolSize; + this.executorService = initExecutorService(); + this.spiderMap = new ConcurrentHashMap(); + } + + public Worker() { + this(DEFAULT_POOL_SIZE); + } + + protected ExecutorService initExecutorService() { + return ThreadUtils.newFixedThreadPool(poolSize); + } + + public void addSpider(Spider spider) { + spider.setExecutorService(executorService); + spiderMap.put(spider.getUUID(), spider); + } + + public Spider getSpider(String uuid){ + return spiderMap.get(uuid); + } + +} diff --git a/webmagic-avalon/webmagic-worker/src/test/java/us/codecraft/webmagic/worker/WorkerTest.java b/webmagic-avalon/webmagic-worker/src/test/java/us/codecraft/webmagic/worker/WorkerTest.java new file mode 100644 index 00000000..24bca19d --- /dev/null +++ b/webmagic-avalon/webmagic-worker/src/test/java/us/codecraft/webmagic/worker/WorkerTest.java @@ -0,0 +1,27 @@ +package us.codecraft.webmagic.worker; + +import org.junit.Test; +import us.codecraft.webmagic.Site; +import us.codecraft.webmagic.Spider; +import us.codecraft.webmagic.processor.PageProcessor; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.*; + +/** + * @author code4crafter@gmail.com + */ +public class WorkerTest { + + @Test + public void testWorkerAsSpiderContains() throws Exception { + PageProcessor pageProcessor = mock(PageProcessor.class); + Site site = mock(Site.class); + when(pageProcessor.getSite()).thenReturn(site); + when(site.getDomain()).thenReturn("codecraft.us"); + Worker worker = new Worker(); + Spider spider = Spider.create(pageProcessor); + worker.addSpider(spider); + assertThat(worker.getSpider("codecraft.us")).isEqualTo(spider); + } +} diff --git a/webmagic-core/src/main/java/us/codecraft/webmagic/Spider.java b/webmagic-core/src/main/java/us/codecraft/webmagic/Spider.java index 0d52ac1f..6fe2880e 100644 --- a/webmagic-core/src/main/java/us/codecraft/webmagic/Spider.java +++ b/webmagic-core/src/main/java/us/codecraft/webmagic/Spider.java @@ -643,6 +643,11 @@ public class Spider implements Runnable, Task { return uuid; } + public Spider setExecutorService(ExecutorService executorService) { + this.executorService = executorService; + return this; + } + @Override public Site getSite() { return site;