some tests for getting source

thomasballinger · thomasballinger · commit 934c511be5e2 · 2014-11-23T00:02:11.000-05:00
diff --git a/bpython/test/test_repl.py b/bpython/test/test_repl.py
@@ -1,8 +1,12 @@
+import collections
+from itertools import islice
 import os
+import socket
 import sys
 import unittest
-from itertools import islice
+
 from mock import Mock, MagicMock
+
 try:
     from unittest import skip
 except ImportError:
@@ -13,6 +17,7 @@ def skip(f):
 
 from bpython import config, repl, cli, autocomplete
 
+
 def setup_config(conf):
     config_struct = config.Struct()
     config.loadini(config_struct, os.devnull)
@@ -253,6 +258,48 @@ def test_nonexistent_name(self):
         self.assertFalse(self.repl.get_args())
 
 
+class TestGetSource(unittest.TestCase):
+    def setUp(self):
+        self.repl = FakeRepl()
+
+    def set_input_line(self, line):
+        """Set current input line of the test REPL."""
+        self.repl.current_line = line
+        self.repl.cursor_offset = len(line)
+
+    def assert_get_source_error_for_current_function(self, func, msg):
+        self.repl.current_func = func
+        self.assertRaises(repl.SourceNotFound, self.repl.get_source_of_current_name)
+        try:
+            self.repl.get_source_of_current_name()
+        except repl.SourceNotFound as e:
+            self.assertEqual(e.args[0], msg)
+
+    def test_current_function(self):
+        self.set_input_line('INPUTLINE')
+        self.repl.current_func = collections.MutableSet.add
+        self.assertTrue("Add an element." in self.repl.get_source_of_current_name())
+
+        self.assert_get_source_error_for_current_function(
+                collections.defaultdict.copy, "No source code found for INPUTLINE")
+
+        self.assert_get_source_error_for_current_function(
+                collections.defaultdict, "could not find class definition")
+
+        self.assert_get_source_error_for_current_function(
+                [], "No source code found for INPUTLINE")
+
+        self.assert_get_source_error_for_current_function(
+                list.pop, "No source code found for INPUTLINE")
+
+    def test_current_line(self):
+        self.repl.interp.locals['a'] = socket.socket
+        self.set_input_line('a')
+        self.assertTrue('dup(self)' in self.repl.get_source_of_current_name())
+
+#TODO add tests for various failures without using current function
+
+
 class TestRepl(unittest.TestCase):
 
     def setInputLine(self, line):
