[SPARK-10417][SQL] Change error message and use assertRaises for unit…

… test
apache · Sep 2, 2015 · f041635 · f041635
1 parent ea2e9d4
commit f041635
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 7 deletions.
diff --git a/python/pyspark/sql/column.py b/python/pyspark/sql/column.py
@@ -227,8 +227,7 @@ def __getattr__(self, item):
         return self.getField(item)
 
     def __iter__(self):
-        raise TypeError('%s.%s object is not iterable' % (self.__module__,
-                                                          self.__class__.__name__))
+        raise TypeError("Column is not iterable")
 
     # string methods
     rlike = _bin_op("rlike")

diff --git a/python/pyspark/sql/tests.py b/python/pyspark/sql/tests.py
@@ -1068,13 +1068,12 @@ def test_with_column_with_existing_name(self):
 
     # regression test for SPARK-10417
     def test_column_iterator(self):
-        # Catch exception raised during improper construction
-        try:
+
+        def foo():
             for x in self.df.key:
                 break
-            self.assertEqual(0, 1)
-        except TypeError:
-            self.assertEqual(1, 1)
+
+        self.assertRaises(TypeError, foo)
 
 
 class HiveContextSQLTests(ReusedPySparkTestCase):