"""Tests for vumi.persist.riak_manager."""

from twisted.trial.unittest import TestCase

from vumi.persist.riak_manager import RiakManager


class TestRiakManager(TestCase):
    pass

"""Tests for vumi.persist.riak_manager."""

from itertools import count

from twisted.trial.unittest import TestCase
from twisted.internet.defer import returnValue

from vumi.persist.riak_manager import RiakManager, flatten_generator
from vumi.persist.tests.test_txriak_manager import CommonRiakManagerTests


class TestRiakManager(CommonRiakManagerTests, TestCase):
    """Most tests are inherited from the CommonRiakManagerTests mixin."""

    def setUp(self):
        self.manager = RiakManager.from_config({'bucket_prefix': 'test.'})
        self.manager.purge_all()

    def tearDown(self):
        self.manager.purge_all()

    def test_call_decorator(self):
        self.assertEqual(RiakManager.call_decorator, flatten_generator)

    def test_flatten_generator(self):
        results = []
        counter = count()

        @flatten_generator
        def f():
            for i in range(3):
                a = yield counter.next()
                results.append(a)

        ret = f()
        self.assertEqual(ret, None)
        self.assertEqual(results, list(range(3)))

    def test_flatter_generator_with_return_value(self):
        @flatten_generator
        def f():
            yield None
            returnValue("foo")

        ret = f()
        self.assertEqual(ret, "foo")

Add tests for (nottx)riak manager.
Add tests for (nottx)riak manager.

  """Tests for vumi.persist.riak_manager."""
  
+ from itertools import count
+ 
  from twisted.trial.unittest import TestCase
+ from twisted.internet.defer import returnValue
  
- from vumi.persist.riak_manager import RiakManager
+ from vumi.persist.riak_manager import RiakManager, flatten_generator
+ from vumi.persist.tests.test_txriak_manager import CommonRiakManagerTests
  
  
- class TestRiakManager(TestCase):
-     pass
+ class TestRiakManager(CommonRiakManagerTests, TestCase):
+     """Most tests are inherited from the CommonRiakManagerTests mixin."""
  
+     def setUp(self):
+         self.manager = RiakManager.from_config({'bucket_prefix': 'test.'})
+         self.manager.purge_all()
+ 
+     def tearDown(self):
+         self.manager.purge_all()
+ 
+     def test_call_decorator(self):
+         self.assertEqual(RiakManager.call_decorator, flatten_generator)
+ 
+     def test_flatten_generator(self):
+         results = []
+         counter = count()
+ 
+         @flatten_generator
+         def f():
+             for i in range(3):
+                 a = yield counter.next()
+                 results.append(a)
+ 
+         ret = f()
+         self.assertEqual(ret, None)
+         self.assertEqual(results, list(range(3)))
+ 
+     def test_flatter_generator_with_return_value(self):
+         @flatten_generator
+         def f():
+             yield None
+             returnValue("foo")
+ 
+         ret = f()
+         self.assertEqual(ret, "foo")
+ 
Add tests for (nottx)riak manager.
## Code Before:
"""Tests for vumi.persist.riak_manager."""

from twisted.trial.unittest import TestCase

from vumi.persist.riak_manager import RiakManager


class TestRiakManager(TestCase):
    pass

## Instruction:
Add tests for (nottx)riak manager.
## Code After:
"""Tests for vumi.persist.riak_manager."""

from itertools import count

from twisted.trial.unittest import TestCase
from twisted.internet.defer import returnValue

from vumi.persist.riak_manager import RiakManager, flatten_generator
from vumi.persist.tests.test_txriak_manager import CommonRiakManagerTests


class TestRiakManager(CommonRiakManagerTests, TestCase):
    """Most tests are inherited from the CommonRiakManagerTests mixin."""

    def setUp(self):
        self.manager = RiakManager.from_config({'bucket_prefix': 'test.'})
        self.manager.purge_all()

    def tearDown(self):
        self.manager.purge_all()

    def test_call_decorator(self):
        self.assertEqual(RiakManager.call_decorator, flatten_generator)

    def test_flatten_generator(self):
        results = []
        counter = count()

        @flatten_generator
        def f():
            for i in range(3):
                a = yield counter.next()
                results.append(a)

        ret = f()
        self.assertEqual(ret, None)
        self.assertEqual(results, list(range(3)))

    def test_flatter_generator_with_return_value(self):
        @flatten_generator
        def f():
            yield None
            returnValue("foo")

        ret = f()
        self.assertEqual(ret, "foo")
