Fixed wait method for callbacks.

[s3rius]

This one makes sure that the processing of messages is over before marking the stream as ended.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 87.04%. Comparing base (74ecf0f) to head (79a0719).
⚠️ Report is 1 commits behind head on master.
✅ All tests successful. No failed tests found.

Additional details and impacted files

@@ Coverage Diff @@## master #50 +/- ##
==========================================
+ Coverage 86.96% 87.04% +0.07% 
==========================================
Files 32 32 Lines 3100 3102 +2 ==========================================
+ Hits 2696 2700 +4 + Misses 404 402 -2 

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.