DataBricks Python unit test error help needed

Question

I have a problem with unit testing in DataBricks. I have not found any similar error message yet. Could someone please help me?

Below you can find my code:

import unittest

def add(n1, n2):
  a = n1+n2
 
  return a

class addTest(unittest.TestCase):
  
  def test_add(self):
    self.assertEqual(add(5,5), 10)

unittest.main()

The error message is this:

EEEEEEEE
======================================================================
ERROR: 36941 (unittest.loader._FailedTest)
----------------------------------------------------------------------
AttributeError: module '__main__' has no attribute '36941'

======================================================================
ERROR: 0 (unittest.loader._FailedTest)
----------------------------------------------------------------------
AttributeError: module '__main__' has no attribute '0'

======================================================================
ERROR: 50000 (unittest.loader._FailedTest)
----------------------------------------------------------------------
AttributeError: module '__main__' has no attribute '50000'

======================================================================
ERROR: 923 (unittest.loader._FailedTest)
----------------------------------------------------------------------
AttributeError: module '__main__' has no attribute '923'

======================================================================
ERROR: dcc9864f9b334f05b20e25d63062f770 (unittest.loader._FailedTest)
----------------------------------------------------------------------
AttributeError: module '__main__' has no attribute 'dcc9864f9b334f05b20e25d63062f770'

======================================================================
ERROR: 3 (unittest.loader._FailedTest)
----------------------------------------------------------------------
AttributeError: module '__main__' has no attribute '3'

======================================================================
ERROR: 4b8784f2c9e2dd8ccd5890b026495bf12fc0671c88baa508721606e240ca39b2 (unittest.loader._FailedTest)
----------------------------------------------------------------------
AttributeError: module '__main__' has no attribute '4b8784f2c9e2dd8ccd5890b026495bf12fc0671c88baa508721606e240ca39b2'

======================================================================
ERROR: unpinned (unittest.loader._FailedTest)
----------------------------------------------------------------------
AttributeError: module '__main__' has no attribute 'unpinned'

Can someone please help me? Thank you!

Answer 1

The unittest.main package requires special handling when using Databricks. The main requirement for use with Databricks is to set exit = False in the list of arguments of unittest.main function. It also makes sense to explicitly pass argv as single-element list, to avoid use of sys.argv that on Databricks contains parameters that were used to start Python subprocess. (see documentation on unittest.main).

But it also may not work as desired. I'm usually generating the test suite by doing autodiscovery myself and the executing it. For example with following code:

import unittest

def discover_test_cases(*test_classes):
  suite = unittest.TestSuite()
  for test_class in test_classes:
    for test in unittest.defaultTestLoader.getTestCaseNames(test_class):
      suite.addTest(test_class(test))
      
  return suite

def discover_test_classes():
  classes = [obj for name, obj in globals().items()
    if name.endswith('Test') and 
       obj.__module__ == '__main__' and 
       isinstance(obj, type) and 
       unittest.case.TestCase in set(obj.__bases__)]

  return discover_test_cases(*classes)

and it's used as:

suite = discover_test_classes()
runner = unittest.TextTestRunner()
results = runner.run(suite)

(you can also use unittest-xml-reporting package and then you can get test results in the JUnit XML format that could be intergrated into build process)

P.S. There is also the nutter library from the Microsoft - it's designed specifically for notebooks testing. It has some advantages, like, automatic discovery of tests, generation of results as JUnit XML, etc. I have an example of using it with Databricks Repos functionality.

I'm using following code