Add an end method to consumers

src/consumer.rs

@@ -32,6 +32,10 @@
 //!        ConsumerState::ConsumerDone
 //!      }
 //!    }
+//!
+//!    fn end(&mut self) {
+//!      println!("finished");
+//!    }
 //!  }
 //!
 //!  // It can consume data directly from a producer
@@ -65,23 +69,28 @@ pub enum ConsumerState {
 /// The run function takes care of continuing or not
 pub trait Consumer {
   fn consume(&mut self, input: &[u8]) -> ConsumerState;
+  fn end(&mut self);
+
   fn run(&mut self, producer: &mut Producer) {
     let mut acc: Vec<u8> = Vec::new();
     //let mut v2: Vec<u8>  = Vec::new();
     loop {
-      let state = producer.produce();
+      let state   = producer.produce();
+      let mut eof = false;
       match state {
         Data(v) => {
           println!("got data");
-          acc.push_all(v)
+          acc.push_all(v);
         },
         Eof([])  => {
           println!("eof empty");
+          eof = true;
           break;
         }
         Eof(v) => {
           println!("eof with {} bytes", v.len());
-          acc.push_all(v)
+          eof = true;
+          acc.push_all(v);
         }
         _ => {break;}
       }
@@ -94,6 +103,7 @@ pub trait Consumer {
         ConsumerDone => {
           println!("data, done");
           acc.clear();
+          self.end();
           //acc.push_all(i);
           break;
         },
@@ -103,6 +113,7 @@ pub trait Consumer {
           //acc.push_all(i);
         }
       }
+    if eof { self.end(); }
     }
   }
 }
@@ -115,12 +126,13 @@ mod tests {
   use std::str;
 
   struct TestPrintConsumer {
-    counter: usize
+    counter: usize,
+    ended:   bool
   }
 
   impl TestPrintConsumer {
     fn new() -> TestPrintConsumer {
-      TestPrintConsumer { counter: 0 }
+      TestPrintConsumer { counter: 0, ended: false }
     }
   }
 
@@ -128,12 +140,17 @@ mod tests {
     fn consume(&mut self, input: &[u8]) -> ConsumerState {
       println!("{} -> {}", self.counter, str::from_utf8(input).unwrap());
       self.counter = self.counter + 1;
-      if self.counter <=4 {
+      if self.counter <= 4 {
         Await
       } else {
         ConsumerDone
       }
     }
+
+    fn end(&mut self) {
+      assert_eq!(self.counter, 5);
+      self.ended = true;
+    }
   }
 
   #[test]
@@ -141,5 +158,7 @@ mod tests {
     let mut p = MemProducer::new("abcdefghijklmnopqrstuvwx".as_bytes(), 4);
     let mut c = TestPrintConsumer::new();
     c.run(&mut p);
+
+    assert!(c.ended);
   }
 }