generalize tuple capturing code so it can be composed

Nathan Marz · Nathan Marz · commit afe8ee8ff454 · 2012-01-25T20:08:26.000-08:00
diff --git a/src/clj/backtype/storm/testing.clj b/src/clj/backtype/storm/testing.clj
@@ -348,26 +348,49 @@
         .get_spout_object
         deserialized-component-object)))
 
+(defn capture-topology [topology]
+  (let [topology (.deepCopy topology)
+        spouts (.get_spouts topology)
+        bolts (.get_bolts topology)
+        all-streams (apply concat
+                           (for [[id spec] (merge (clojurify-structure spouts)
+                                                  (clojurify-structure bolts))]
+                             (for [[stream info] (.. spec get_common get_streams)]
+                               [(GlobalStreamId. id stream) (.is_direct info)])))
+        capturer (TupleCaptureBolt.)]
+    (.set_bolts topology
+                (assoc (clojurify-structure bolts)
+                  (uuid)
+                  (Bolt.                   
+                   (serialize-component-object capturer)
+                   (mk-plain-component-common (into {} (for [[id direct?] all-streams]
+                                                         [id (if direct?
+                                                               (mk-direct-grouping)
+                                                               (mk-global-grouping))]))
+                                              {}
+                                              nil))
+                  ))
+    {:topology topology
+     :capturer capturer}
+    ))
+
 ;; TODO: mock-sources needs to be able to mock out state spouts as well
 (defnk complete-topology [cluster-map topology :mock-sources {} :storm-conf {}]
   ;; TODO: the idea of mocking for transactional topologies should be done an
   ;; abstraction level above... should have a complete-transactional-topology for this
-  (let [storm-name (str "topologytest-" (uuid))
+  (let [{topology :topology capturer :capturer} (capture-topology topology)
+        storm-name (str "topologytest-" (uuid))
         state (:storm-cluster-state cluster-map)
         spouts (.get_spouts topology)
-        bolts (.get_bolts topology)
         replacements (map-val (fn [v]
                                 (FixedTupleSpout.
                                  (for [tup v]
                                    (if (map? tup)
                                      (FixedTuple. (:stream tup) (:values tup))
                                      tup))))
                               mock-sources)
-        all-streams (apply concat
-                           (for [[id spec] (merge (clojurify-structure spouts) (clojurify-structure bolts))]
-                             (for [[stream info] (.. spec get_common get_streams)]
-                               [(GlobalStreamId. id stream) (.is_direct info)])))
-        capturer (TupleCaptureBolt. storm-name)
+        
+
         ]
     (doseq [[id spout] replacements]
       (let [spout-spec (get spouts id)]
@@ -378,18 +401,7 @@
         (throw (RuntimeException. "Cannot complete topology unless every spout is a CompletableSpout (or mocked to be)"))
         ))
     
-    (.set_bolts topology
-                (assoc (clojurify-structure bolts)
-                  (uuid)
-                  (Bolt.                   
-                   (serialize-component-object capturer)
-                   (mk-plain-component-common (into {} (for [[id direct?] all-streams]
-                                                         [id (if direct?
-                                                               (mk-direct-grouping)
-                                                               (mk-global-grouping))]))
-                                              {}
-                                              nil))
-                  ))
+    
     (submit-local-topology (:nimbus cluster-map) storm-name storm-conf topology)
     
     
@@ -403,7 +415,7 @@
       (doseq [spout (spout-objects spouts)]
         (cleanup spout)))
 
-    (.getResults capturer)
+    (.getAndClearResults capturer)
     ))
 
 (defn read-tuples
diff --git a/src/jvm/backtype/storm/testing/TupleCaptureBolt.java b/src/jvm/backtype/storm/testing/TupleCaptureBolt.java
@@ -8,18 +8,18 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.UUID;
 
 
 public class TupleCaptureBolt implements IBolt {
     public static transient Map<String, Map<String, List<FixedTuple>>> emitted_tuples = new HashMap<String, Map<String, List<FixedTuple>>>();
 
     private String _name;
-    private Map<String, List<FixedTuple>> _results = null;
     private OutputCollector _collector;
 
-    public TupleCaptureBolt(String name) {
-        _name = name;
-        emitted_tuples.put(name, new HashMap<String, List<FixedTuple>>());
+    public TupleCaptureBolt() {
+        _name = UUID.randomUUID().toString();
+        emitted_tuples.put(_name, new HashMap<String, List<FixedTuple>>());
     }
 
     public void prepare(Map stormConf, TopologyContext context, OutputCollector collector) {
@@ -37,13 +37,14 @@ public void execute(Tuple input) {
     }
 
     public Map<String, List<FixedTuple>> getResults() {
-        if(_results==null) {
-            _results = emitted_tuples.remove(_name);
-        }
-        return _results;
+        return emitted_tuples.get(_name);
     }
 
     public void cleanup() {
     }
+    
+    public Map<String, List<FixedTuple>> getAndClearResults() {
+        return emitted_tuples.remove(_name);
+    }
 
 }
diff --git a/test/clj/backtype/storm/transactional_test.clj b/test/clj/backtype/storm/transactional_test.clj
@@ -317,6 +317,16 @@
          (.setBolt "count" (KeyedSummingBatchBolt.) 2)
          (.fieldsGrouping "spout" (Fields. ["word"])))
 
+     ;; make a complex topology, and then make a final bolt at the end that doesn't ack (and lets me manipulate its tuples from here)
+
+     ;; let it go to completion
+     ;; make sure there are no commits
+     ;; ack tuples for txid = 2
+     ;; make sure there are no commits
+     ;; ack tuples for txid = 1, let it run
+     ;; make sure there are 2 commits + 2 more batch emits
+     ;; extract the code to add a capturing bolt and transform the topology
+     
      (add-transactional-data data
                              {0 [["dog" 3]
                                  ["cat" 4]
