from time import sleep
import unittest

from tweepy.api import API
from tweepy.streaming import Stream, StreamListener

from config import create_auth
from mock import mock_tweet

class MockStreamListener(StreamListener):
    def __init__(self):
        super(MockStreamListener, self).__init__()
        self.status_count = 0

    def on_status(self, status):
        self.status_count += 1
        return False

class TweepyStreamTests(unittest.TestCase):
    def setUp(self):
        self.auth = create_auth()
        self.listener = MockStreamListener()
        self.stream = Stream(self.auth, self.listener)

    def tearDown(self):
        self.stream.disconnect()

    def test_userstream(self):
        self.stream.userstream(async=True)

        # Generate random tweet which should show up in the stream.
        # Wait a bit of time for it to arrive before asserting.
        API(self.auth).update_status(mock_tweet())
        sleep(1)

        self.assertEqual(self.listener.status_count, 1)


@@ -60,16 +60,49 @@
 ort API%0A
+from tweepy.models import Status%0A
 from twe
@@ -265,24 +265,35 @@
 _init__(self
+, test_case
 ):%0A        s
@@ -351,90 +351,392 @@
 elf.
-status_count = 0%0A%0A    def on_status(self, status):%0A        self.status_count += 1%0A
+test_case = test_case%0A        self.status_count = 0%0A        self.status_stop_count = 0%0A        self.connect_cb = None%0A%0A    def on_connect(self):%0A        if self.connect_cb:%0A            self.connect_cb()%0A%0A    def on_status(self, status):%0A        self.status_count += 1%0A        self.test_case.assertIsInstance(status, Status)%0A        if self.status_stop_count == self.status_count:%0A    
     
@@ -895,16 +895,20 @@
 istener(
+self
 )%0A      
@@ -1050,52 +1050,8 @@
 f):%0A
-        self.stream.userstream(async=True)%0A%0A
     
@@ -1126,63 +1126,204 @@
     
-# Wait a bit of time for it to arrive before asserting.
+def on_connect():%0A            API(self.auth).update_status(mock_tweet())%0A%0A        self.listener.connect_cb = on_connect%0A        self.listener.status_stop_count = 1%0A        self.stream.userstream()
 %0A   
@@ -1331,68 +1331,157 @@
     
-API(
 self.a
-uth).update_status(mock_tweet())%0A        sleep(1)%0A
+ssertEqual(self.listener.status_count, 1)%0A%0A    def test_sample(self):%0A        self.listener.status_stop_count = 10%0A        self.stream.sample()
 %0A   
@@ -1493,32 +1493,33 @@
 self.assertEqual
+s
 (self.listener.s
@@ -1514,29 +1514,85 @@
 f.listener.status_count,
- 1
+%0A                          self.listener.status_stop_count
 )%0A%0A

Add streaming test_sample test.