fix streaming tests

1 files changed, 9 insertions, 4 deletions

diff --git a/tests/test_streaming.py b/tests/test_streaming.py
index 1f8d75f..33a3f9a 100644
--- a/tests/test_streaming.py
+++ b/tests/test_streaming.py
@@ -1,7 +1,9 @@
 import six
 import pytest
 import itertools
-from mastodon.streaming import StreamListener, MalformedEventError
+from mastodon.streaming import StreamListener
+from mastodon.Mastodon import MastodonMalformedEventError
+
 
 
 class Listener(StreamListener):
@@ -20,6 +22,9 @@ class Listener(StreamListener):
     def on_delete(self, status_id):
         self.deletes.append(status_id)
 
+    def on_blahblah(self, data):
+        pass
+
     def handle_heartbeat(self):
         self.heartbeats += 1
 
@@ -101,7 +106,7 @@ def test_unknown_event():
 
 def test_missing_event_name():
     listener = Listener()
-    with pytest.raises(MalformedEventError):
+    with pytest.raises(MastodonMalformedEventError):
         listener.handle_stream_([
             'data: {}',
             '',
@@ -115,7 +120,7 @@ def test_missing_event_name():
 
 def test_missing_data():
     listener = Listener()
-    with pytest.raises(MalformedEventError):
+    with pytest.raises(MastodonMalformedEventError):
         listener.handle_stream_([
             'event: update',
             '',
@@ -169,7 +174,7 @@ def test_valid_utf8():
 def test_invalid_utf8():
     """Cat Face With Tears O"""
     listener = Listener()
-    with pytest.raises(MalformedEventError):
+    with pytest.raises(MastodonMalformedEventError):
         listener.handle_stream_([
             'event: update',
             'data: {"foo": "\xF0\x9F\x98"}',