Fix Events::drain bug

Co-authored-by: Alice Cecile <[email protected]>

crates/bevy_ecs/src/event.rs

@@ -336,6 +336,10 @@ impl<T: Component> Events<T> {
 
     /// Creates a draining iterator that removes all events.
     pub fn drain(&mut self) -> impl Iterator<Item = T> + '_ {
+        self.a_start_event_count = 0;
+        self.b_start_event_count = 0;
+        self.event_count = 0;
+
         let map = |i: EventInstance<T>| i.event;
         match self.state {
             State::A => self
@@ -480,4 +484,24 @@ mod tests {
     ) -> Vec<TestEvent> {
         reader.iter(events).cloned().collect::<Vec<TestEvent>>()
     }
+
+    #[test]
+    fn test_events_drain() {
+        struct E;
+        let mut events = Events::<E>::default();
+        let mut reader = events.get_reader();
+
+        assert!(reader.iter(&events).next().is_none());
+
+        events.send(E);
+        assert!(reader.iter(&events).next().is_some());
+        assert!(reader.iter(&events).next().is_none());
+
+        let _ = events.drain();
+        assert!(reader.iter(&events).next().is_none());
+
+        events.send(E);
+        assert!(reader.iter(&events).next().is_some());
+        assert!(reader.iter(&events).next().is_none());
+    }
 }