Add tests for %timeit with AST transformations

IPython/core/tests/test_interactiveshell.py

@@ -438,6 +438,21 b' class TestAstTransform(unittest.TestCase):'
         with tt.AssertNotPrints('-55'):
             ip.run_cell('print (n)')
 
+    def test_timeit(self):
+        called = set()
+        def f(x):
+            called.add(x)
+        ip.push({'f':f})
+        
+        with tt.AssertPrints("best of "):
+            ip.run_line_magic("timeit", "-n1 f(1)")
+        self.assertEqual(called, set([-1]))
+        called.clear()
+        
+        with tt.AssertPrints("best of "):
+            ip.run_cell_magic("timeit", "-n1 f(2)", "f(3)")
+        self.assertEqual(called, set([-2, -3]))
+
 class IntegerWrapper(ast.NodeTransformer):
     """Wraps all integers in a call to Integer()"""
     def visit_Num(self, node):
@@ -453,6 +468,7 b' class TestAstTransform2(unittest.TestCase):'
         self.calls = []
         def Integer(*args):
             self.calls.append(args)
+            return args
         ip.push({"Integer": Integer})
     
     def tearDown(self):
@@ -463,6 +479,21 b' class TestAstTransform2(unittest.TestCase):'
         ip.run_cell("n = 2")
         self.assertEqual(self.calls, [(2,)])
 
+    def test_timeit(self):
+        called = set()
+        def f(x):
+            called.add(x)
+        ip.push({'f':f})
+        
+        with tt.AssertPrints("best of "):
+            ip.run_line_magic("timeit", "-n1 f(1)")
+        self.assertEqual(called, set([(1,)]))
+        called.clear()
+        
+        with tt.AssertPrints("best of "):
+            ip.run_cell_magic("timeit", "-n1 f(2)", "f(3)")
+        self.assertEqual(called, set([(2,), (3,)]))
+
 class ErrorTransformer(ast.NodeTransformer):
     """Throws an error when it sees a number."""
     def visit_Num(self):