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.