1 files changed, 28 insertions, 15 deletions
diff --git a/Lib/test/test_dbm.py b/Lib/test/test_dbm.py
index 26d4c14..d11d924 100644
--- a/Lib/test/test_dbm.py
+++ b/Lib/test/test_dbm.py
@@ -1,4 +1,3 @@
-#! /usr/bin/env python3
 """Test script for the dbm.open function based on testdumbdbm.py"""
 
 import os
@@ -9,6 +8,11 @@ import test.support
 # Skip tests if dbm module doesn't exist.
 dbm = test.support.import_module('dbm')
 
+try:
+    from dbm import ndbm
+except ImportError:
+    ndbm = None
+
 _fname = test.support.TESTFN
 
 #
@@ -34,7 +38,7 @@ def delete_files():
         test.support.unlink(f)
 
 
-class AnyDBMTestCase(unittest.TestCase):
+class AnyDBMTestCase:
     _dict = {'0': b'',
              'a': b'Python:',
              'b': b'Programming',
@@ -71,8 +75,8 @@ class AnyDBMTestCase(unittest.TestCase):
         f.close()
 
     def test_anydbm_creation_n_file_exists_with_invalid_contents(self):
-        with open(_fname, "w") as w:
-            pass # create an empty file
+        # create an empty file
+        test.support.create_empty_file(_fname)
 
         f = dbm.open(_fname, 'n')
         self.addCleanup(f.close)
@@ -119,10 +123,6 @@ class AnyDBMTestCase(unittest.TestCase):
 
 
 class WhichDBTestCase(unittest.TestCase):
-    # Actual test methods are added to namespace after class definition.
-    def __init__(self, *args):
-        unittest.TestCase.__init__(self, *args)
-
     def test_whichdb(self):
         for module in dbm_iterator():
             # Check whether whichdb correctly guesses module name
@@ -134,7 +134,7 @@ class WhichDBTestCase(unittest.TestCase):
             delete_files()
             f = module.open(_fname, 'c')
             f.close()
-            self.assertEqual(name, dbm.whichdb(_fname))
+            self.assertEqual(name, self.dbm.whichdb(_fname))
             # Now add a key
             f = module.open(_fname, 'w')
             f[b"1"] = b"1"
@@ -143,7 +143,15 @@ class WhichDBTestCase(unittest.TestCase):
             # and read it
             self.assertTrue(f[b"1"] == b"1")
             f.close()
-            self.assertEqual(name, dbm.whichdb(_fname))
+            self.assertEqual(name, self.dbm.whichdb(_fname))
+
+    @unittest.skipUnless(ndbm, reason='Test requires ndbm')
+    def test_whichdb_ndbm(self):
+        # Issue 17198: check that ndbm which is referenced in whichdb is defined
+        db_file = '{}_ndbm.db'.format(_fname)
+        with open(db_file, 'w'):
+            self.addCleanup(test.support.unlink, db_file)
+        self.assertIsNone(self.dbm.whichdb(db_file[:-3]))
 
     def tearDown(self):
         delete_files()
@@ -153,6 +161,7 @@ class WhichDBTestCase(unittest.TestCase):
         self.filename = test.support.TESTFN
         self.d = dbm.open(self.filename, 'c')
         self.d.close()
+        self.dbm = test.support.import_fresh_module('dbm')
 
     def test_keys(self):
         self.d = dbm.open(self.filename, 'c')
@@ -169,12 +178,16 @@ class WhichDBTestCase(unittest.TestCase):
         self.d.close()
 
 
-def test_main():
-    classes = [WhichDBTestCase]
+def load_tests(loader, tests, pattern):
+    classes = []
     for mod in dbm_iterator():
-        classes.append(type("TestCase-" + mod.__name__, (AnyDBMTestCase,),
+        classes.append(type("TestCase-" + mod.__name__,
+                            (AnyDBMTestCase, unittest.TestCase),
                             {'module': mod}))
-    test.support.run_unittest(*classes)
+    suites = [unittest.makeSuite(c) for c in classes]
+
+    tests.addTests(suites)
+    return tests
 
 if __name__ == "__main__":
-    test_main()
+    unittest.main()