from ..utils import TranspileTestCase


class ReturnTests(TranspileTestCase):
    def test_return_methods(self):
        # Empty return
        self.assertCodeExecution("""
            def m():
                print('haha')
                return
                print('haha')

            print(m())
        """)

        # Single return
        self.assertCodeExecution("""
            def m():
                print('haha')
                return 1
                print('haha')

            print(m() + 1)
        """)

        # Multiple return
        self.assertCodeExecution("""
            def m():
                print('haha')
                return 1, 2
                print('haha')

            print(m())
        """)

@@ -202,36 +202,45 @@
          print('
-haha
+before return
 ')%0A             
@@ -264,36 +264,44 @@
          print('
-haha
+after return
 ')%0A%0A            
@@ -422,36 +422,45 @@
          print('
-haha
+before return
 ')%0A             
@@ -486,36 +486,44 @@
          print('
-haha
+after return
 ')%0A%0A            
@@ -658,20 +658,29 @@
  print('
-haha
+before return
 ')%0A     
@@ -729,12 +729,20 @@
 nt('
-haha
+after return
 ')%0A%0A

use more informatic print's in test_return.py