testing progress

Nathan Marz · Nathan Marz · commit 26cc00aa45c3 · 2012-01-25T15:33:12.000-08:00
diff --git a/src/jvm/backtype/storm/testing/KeyedCountingBatchBolt.java b/src/jvm/backtype/storm/testing/KeyedCountingBatchBolt.java
@@ -0,0 +1,44 @@
+package backtype.storm.testing;
+
+import backtype.storm.task.TopologyContext;
+import backtype.storm.topology.OutputFieldsDeclarer;
+import backtype.storm.topology.base.BaseBatchBolt;
+import backtype.storm.coordination.BatchOutputCollector;
+import backtype.storm.tuple.Fields;
+import backtype.storm.tuple.Tuple;
+import backtype.storm.tuple.Values;
+import backtype.storm.utils.Utils;
+import java.util.HashMap;
+import java.util.Map;
+
+public class KeyedCountingBatchBolt extends BaseBatchBolt {
+    BatchOutputCollector _collector;
+    Object _id;
+    Map<Object, Integer> _counts = new HashMap<Object, Integer>();
+    
+    @Override
+    public void prepare(Map conf, TopologyContext context, BatchOutputCollector collector, Object id) {
+        _collector = collector;
+        _id = id;
+    }
+
+    @Override
+    public void execute(Tuple tuple) {
+        Object key = tuple.getValue(1);
+        int curr = Utils.get(_counts, key, 0);
+        _counts.put(key, curr + 1);
+    }
+
+    @Override
+    public void finishBatch() {
+        for(Object key: _counts.keySet()) {
+            _collector.emit(new Values(_id, key, _counts.get(key)));
+        }
+    }   
+
+    @Override
+    public void declareOutputFields(OutputFieldsDeclarer declarer) {
+        declarer.declare(new Fields("tx", "key", "count"));
+    }
+    
+}
diff --git a/src/jvm/backtype/storm/testing/KeyedCountingCommitterBolt.java b/src/jvm/backtype/storm/testing/KeyedCountingCommitterBolt.java
@@ -0,0 +1,43 @@
+package backtype.storm.testing;
+
+import backtype.storm.task.TopologyContext;
+import backtype.storm.topology.OutputFieldsDeclarer;
+import backtype.storm.coordination.BatchOutputCollector;
+import backtype.storm.topology.base.BaseCommitterBolt;
+import backtype.storm.transactional.TransactionAttempt;
+import backtype.storm.tuple.Fields;
+import backtype.storm.tuple.Tuple;
+import backtype.storm.tuple.Values;
+import backtype.storm.utils.Utils;
+import java.util.HashMap;
+import java.util.Map;
+
+public class KeyedCountingCommitterBolt extends BaseCommitterBolt {
+    TransactionAttempt _id;
+    Map<Object, Integer> _counts = new HashMap<Object, Integer>();
+    
+    @Override
+    public void prepare(Map conf, TopologyContext context, TransactionAttempt id) {
+        _id = id;
+    }
+
+    @Override
+    public void execute(Tuple tuple) {
+        Object key = tuple.getValue(1);
+        int curr = Utils.get(_counts, key, 0);
+        _counts.put(key, curr + 1);
+    }
+
+    @Override
+    public void commit(BatchOutputCollector collector) {
+        for(Object key: _counts.keySet()) {
+            collector.emit(new Values(_id, key, _counts.get(key)));
+        }
+    }   
+
+    @Override
+    public void declareOutputFields(OutputFieldsDeclarer declarer) {
+        declarer.declare(new Fields("tx", "key", "count"));
+    }
+    
+}
diff --git a/src/jvm/backtype/storm/testing/KeyedSummingBatchBolt.java b/src/jvm/backtype/storm/testing/KeyedSummingBatchBolt.java
@@ -0,0 +1,43 @@
+package backtype.storm.testing;
+
+import backtype.storm.coordination.BatchOutputCollector;
+import backtype.storm.task.TopologyContext;
+import backtype.storm.topology.OutputFieldsDeclarer;
+import backtype.storm.topology.base.BaseBatchBolt;
+import backtype.storm.tuple.Fields;
+import backtype.storm.tuple.Tuple;
+import backtype.storm.tuple.Values;
+import backtype.storm.utils.Utils;
+import java.util.HashMap;
+import java.util.Map;
+
+public class KeyedSummingBatchBolt extends BaseBatchBolt {
+    BatchOutputCollector _collector;
+    Object _id;
+    Map<Object, Integer> _sums = new HashMap<Object, Integer>();
+    
+    @Override
+    public void prepare(Map conf, TopologyContext context, BatchOutputCollector collector, Object id) {
+        _collector = collector;
+        _id = id;
+    }
+
+    @Override
+    public void execute(Tuple tuple) {
+        Object key = tuple.getValue(1);
+        int curr = Utils.get(_sums, key, 0);
+        _sums.put(key, curr + tuple.getInteger(2));
+    }
+
+    @Override
+    public void finishBatch() {
+        for(Object key: _sums.keySet()) {
+            _collector.emit(new Values(_id, key, _sums.get(key)));
+        }
+    }   
+
+    @Override
+    public void declareOutputFields(OutputFieldsDeclarer declarer) {
+        declarer.declare(new Fields("tx", "key", "sum"));
+    }    
+}
diff --git a/test/clj/backtype/storm/transactional_test.clj b/test/clj/backtype/storm/transactional_test.clj
@@ -4,7 +4,8 @@
   (:import [backtype.storm.transactional TransactionalSpoutCoordinator ITransactionalSpout ITransactionalSpout$Coordinator TransactionAttempt
             TransactionalTopologyBuilder])
   (:import [backtype.storm.transactional.state TransactionalState RotatingTransactionalState RotatingTransactionalState$StateInitializer])
-  (:import [backtype.storm.testing CountingBatchBolt MemoryTransactionalSpout])
+  (:import [backtype.storm.testing CountingBatchBolt MemoryTransactionalSpout
+            KeyedCountingBatchBolt KeyedCountingCommitterBolt KeyedSummingBatchBolt])
   (:use [backtype.storm bootstrap testing])
   (:use [backtype.storm.daemon common])  
   )
@@ -312,7 +313,9 @@
                                                (Fields. ["word" "amt"])
                                                3)
                     2))
-
+     (-> builder
+         (.setBolt "count" (KeyedSummingBatchBolt.) 2)
+         (.fieldsGrouping "spout" (Fields. ["word"])))
 
      (add-transactional-data data
                              {0 [["dog" 3]
@@ -329,6 +332,6 @@
            (complete-topology cluster
                               (.buildTopology builder)
                               :storm-conf {TOPOLOGY-DEBUG true}))
-     (println results)
+     (println (read-tuples results "count"))
 
      )))
