sdks/python/apache_beam/dataframe/doctests_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import doctest
import os
import tempfile
import unittest

from apache_beam.dataframe import doctests

SAMPLE_DOCTEST = '''
>>> df = pd.DataFrame({'Animal': ['Falcon', 'Falcon',
...                               'Parrot', 'Parrot'],
...                    'Max Speed': [380., 370., 24., 26.]})
>>> df
   Animal  Max Speed
0  Falcon      380.0
1  Falcon      370.0
2  Parrot       24.0
3  Parrot       26.0
>>> df.groupby(['Animal']).mean()
        Max Speed
Animal
Falcon      375.0
Parrot       25.0
'''

CHECK_USES_DEFERRED_DATAFRAMES = '''
>>> type(pd).__name__
'FakePandasObject'

>>> type(pd.DataFrame([]))
<class 'apache_beam.dataframe.frames.DeferredDataFrame'>

>>> type(pd.DataFrame.from_dict({'a': [1, 2], 'b': [3, 4]}))
<class 'apache_beam.dataframe.frames.DeferredDataFrame'>

>>> pd.Index(range(10))
RangeIndex(start=0, stop=10, step=1)
'''

WONT_IMPLEMENT_RAISING_TESTS = '''
>>> import apache_beam
>>> raise apache_beam.dataframe.frame_base.WontImplementError('anything')
ignored exception
>>> pd.Series(range(10)).__array__()
ignored result
'''

ERROR_RAISING_NAME_ERROR_TESTS = '''
>>> import apache_beam
>>> raise %s('anything')
ignored exception
>>> raise NameError
ignored exception
>>> undefined_name
ignored exception
>>> 2 + 2
4
>>> raise NameError
failed exception
'''

WONT_IMPLEMENT_RAISING_NAME_ERROR_TESTS = ERROR_RAISING_NAME_ERROR_TESTS % (
    'apache_beam.dataframe.frame_base.WontImplementError', )

NOT_IMPLEMENTED_RAISING_TESTS = '''
>>> import apache_beam
>>> raise NotImplementedError('anything')
ignored exception
'''

NOT_IMPLEMENTED_RAISING_NAME_ERROR_TESTS = ERROR_RAISING_NAME_ERROR_TESTS % (
    'NotImplementedError', )

FAILED_ASSIGNMENT = '''
>>> def foo(): raise NotImplementedError()
>>> res = 'old_value'
>>> res = foo()
>>> print(res)
ignored NameError
'''

RST_IPYTHON = '''
Here is an example
.. ipython::

    2 + 2

some multi-line examples

.. ipython::

    def foo(x):
        return x * x
    foo(4)
    foo(
        4
    )

    In [100]: def foo(x):
       ....:     return x * x * x
       ....:
    foo(5)

history is preserved

    foo(3)
    foo(4)

and finally an example with pandas

.. ipython::

    pd.Series([1, 2, 3]).max()


This one should be skipped:

.. ipython::

   @verbatim
   not run or tested

and someting that'll fail (due to fake vs. real pandas)

.. ipython::

   type(pd)
'''


class DoctestTest(unittest.TestCase):
  def test_good(self):
    result = doctests.teststring(SAMPLE_DOCTEST, report=False)
    self.assertEqual(result.attempted, 3)
    self.assertEqual(result.failed, 0)

  def test_failure(self):
    result = doctests.teststring(
        SAMPLE_DOCTEST.replace('25.0', '25.00001'), report=False)
    self.assertEqual(result.attempted, 3)
    self.assertEqual(result.failed, 1)

  def test_uses_beam_dataframes(self):
    result = doctests.teststring(CHECK_USES_DEFERRED_DATAFRAMES, report=False)
    self.assertNotEqual(result.attempted, 0)
    self.assertEqual(result.failed, 0)

  def test_file(self):
    with tempfile.TemporaryDirectory() as dir:
      filename = os.path.join(dir, 'tests.py')
      with open(filename, 'w') as fout:
        fout.write(SAMPLE_DOCTEST)
      result = doctests.testfile(filename, module_relative=False, report=False)
    self.assertEqual(result.attempted, 3)
    self.assertEqual(result.failed, 0)

  def test_file_uses_beam_dataframes(self):
    with tempfile.TemporaryDirectory() as dir:
      filename = os.path.join(dir, 'tests.py')
      with open(filename, 'w') as fout:
        fout.write(CHECK_USES_DEFERRED_DATAFRAMES)
      result = doctests.testfile(filename, module_relative=False, report=False)
    self.assertNotEqual(result.attempted, 0)
    self.assertEqual(result.failed, 0)

  def test_wont_implement(self):
    result = doctests.teststring(
        WONT_IMPLEMENT_RAISING_TESTS,
        optionflags=doctest.ELLIPSIS,
        wont_implement_ok=True)
    self.assertNotEqual(result.attempted, 0)
    self.assertEqual(result.failed, 0)

    result = doctests.teststring(
        WONT_IMPLEMENT_RAISING_TESTS,
        optionflags=doctest.IGNORE_EXCEPTION_DETAIL,
        wont_implement_ok=True)
    self.assertNotEqual(result.attempted, 0)
    self.assertEqual(result.failed, 0)

  def test_wont_implement_followed_by_name_error(self):
    result = doctests.teststring(
        WONT_IMPLEMENT_RAISING_NAME_ERROR_TESTS,
        optionflags=doctest.ELLIPSIS,
        wont_implement_ok=True)
    self.assertEqual(result.attempted, 6)
    self.assertEqual(result.failed, 1)  # Only the very last one.

  def test_not_implemented(self):
    result = doctests.teststring(
        NOT_IMPLEMENTED_RAISING_TESTS,
        optionflags=doctest.ELLIPSIS,
        not_implemented_ok=True)
    self.assertNotEqual(result.attempted, 0)
    self.assertEqual(result.failed, 0)

    result = doctests.teststring(
        NOT_IMPLEMENTED_RAISING_TESTS,
        optionflags=doctest.IGNORE_EXCEPTION_DETAIL,
        not_implemented_ok=True)
    self.assertNotEqual(result.attempted, 0)
    self.assertEqual(result.failed, 0)

  def test_not_implemented_followed_by_name_error(self):
    result = doctests.teststring(
        NOT_IMPLEMENTED_RAISING_NAME_ERROR_TESTS,
        optionflags=doctest.ELLIPSIS,
        not_implemented_ok=True)
    self.assertEqual(result.attempted, 6)
    self.assertEqual(result.failed, 1)  # Only the very last one.

  def test_failed_assignment(self):
    result = doctests.teststring(
        FAILED_ASSIGNMENT,
        optionflags=doctest.ELLIPSIS,
        not_implemented_ok=True)
    self.assertNotEqual(result.attempted, 0)
    self.assertEqual(result.failed, 0)

  def test_rst_ipython(self):
    try:
      import IPython
    except ImportError:
      raise unittest.SkipTest('IPython not available')
    result = doctests.test_rst_ipython(RST_IPYTHON, 'test_rst_ipython')
    self.assertEqual(result.attempted, 8)
    self.assertEqual(result.failed, 1)  # Only the very last one.


if __name__ == '__main__':
  unittest.main()