From bba7d7184d4e07dc5123c44c82103d496b9058c2 Mon Sep 17 00:00:00 2001
From: chaos <eagleqingluo@gmail.com>
Date: Tue, 19 Sep 2017 12:07:26 +0800
Subject: [PATCH] feat(RedisScheduler) add flush function to clean redis data

add flush function to clean redis data
---
 .../webmagic/scheduler/RedisScheduler.java    | 15 ++++++++++
 .../scheduler/RedisSchedulerTest.java         | 28 +++++++++++++++----
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/webmagic-extension/src/main/java/us/codecraft/webmagic/scheduler/RedisScheduler.java b/webmagic-extension/src/main/java/us/codecraft/webmagic/scheduler/RedisScheduler.java
index ce1111f24..1090c9aca 100644
--- a/webmagic-extension/src/main/java/us/codecraft/webmagic/scheduler/RedisScheduler.java
+++ b/webmagic-extension/src/main/java/us/codecraft/webmagic/scheduler/RedisScheduler.java
@@ -126,4 +126,19 @@ public int getTotalRequestsCount(Task task) {
             pool.returnResource(jedis);
         }
     }
+
+    /**
+     * clean all current task redis queue
+     * @param task current task
+     */
+    public void flush(Task task) {
+        Jedis jedis = pool.getResource();
+        try {
+            jedis.del(getSetKey(task));
+            jedis.del(getQueueKey(task));
+            jedis.del(getItemKey(task));
+        } finally {
+            pool.returnResource(jedis);
+        }
+    }
 }
diff --git a/webmagic-extension/src/test/java/us/codecraft/webmagic/scheduler/RedisSchedulerTest.java b/webmagic-extension/src/test/java/us/codecraft/webmagic/scheduler/RedisSchedulerTest.java
index b4124d2d9..541367c6f 100644
--- a/webmagic-extension/src/test/java/us/codecraft/webmagic/scheduler/RedisSchedulerTest.java
+++ b/webmagic-extension/src/test/java/us/codecraft/webmagic/scheduler/RedisSchedulerTest.java
@@ -15,16 +15,13 @@
 public class RedisSchedulerTest {
 
     private RedisScheduler redisScheduler;
+    private Task task;
 
     @Before
     public void setUp() {
         redisScheduler = new RedisScheduler("localhost");
-    }
 
-    @Ignore("environment depended")
-    @Test
-    public void test() {
-        Task task = new Task() {
+        task = new Task() {
             @Override
             public String getUUID() {
                 return "1";
@@ -35,6 +32,11 @@ public Site getSite() {
                 return null;
             }
         };
+    }
+
+    @Ignore("environment depended")
+    @Test
+    public void test() {
         Request request = new Request("http://www.ibm.com/developerworks/cn/java/j-javadev2-22/");
         request.putExtra("1","2");
         redisScheduler.push(request, task);
@@ -42,4 +44,20 @@ public Site getSite() {
         assertThat(poll).isEqualTo(request);
 
     }
+
+    @Ignore("environment depended")
+    @Test
+    public void testFlush() {
+        Request request = new Request("http://www.baidu.com");
+        request.putExtra("1","2");
+        redisScheduler.push(request, task);
+        int totalRequestsCount = redisScheduler.getTotalRequestsCount(task);
+        assertThat(totalRequestsCount).isGreaterThan(0);
+        redisScheduler.flush(task);
+        totalRequestsCount = redisScheduler.getTotalRequestsCount(task);
+        assertThat(totalRequestsCount).isEqualTo(0);
+        int leftRequestsCount = redisScheduler.getLeftRequestsCount(task);
+        assertThat(leftRequestsCount).isEqualTo(0);
+
+    }
 }