Include pytest and mock as test requirements, together with a test

command to run the tests using "python setup.py test".

Author: dukebody

setup.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
 from setuptools import setup
+from setuptools.command.test import test as TestCommand
 import re
 
 for line in open('sklearn_pandas/__init__.py'):
@@ -9,6 +10,24 @@
         __version__, = match.groups()
 
 
+class PyTest(TestCommand):
+    user_options = [('pytest-args=', 'a', "Arguments to pass to py.test")]
+
+    def initialize_options(self):
+        TestCommand.initialize_options(self)
+        self.pytest_args = []
+
+    def finalize_options(self):
+        TestCommand.finalize_options(self)
+        self.test_args = []
+        self.test_suite = True
+
+    def run(self):
+        import pytest
+        errno = pytest.main(self.pytest_args)
+        raise SystemExit(errno)
+
+
 setup(name='sklearn-pandas',
       version=__version__,
       description='Pandas integration with sklearn',
@@ -21,5 +40,7 @@
           'scikit-learn>=0.13',
           'scipy>=0.14',
           'pandas>=0.11.0',
-          'numpy>=1.6.1']
-)
+          'numpy>=1.6.1'],
+      tests_require=['pytest', 'mock'],
+      cmdclass={'test': PyTest},
+      )

tests/test_dataframe_mapper.py

@@ -97,6 +97,10 @@ def test_cols_list_column_vector():
 
 
 def test_list_transformers():
+    """
+    Specifying a list of transformers applies them sequentially to the
+    selected column.
+    """
     dataframe = pd.DataFrame({"a": [1, np.nan, 3], "b": [1, 5, 7]})
 
     mapper = DataFrameMapper([