logging test less brittle

Author: mitzimorris

test/test_sample.py

@@ -1801,48 +1801,38 @@ def test_diagnostics(self):
         stan = os.path.join(DATAFILES_PATH, 'eight_schools.stan')
         model = CmdStanModel(stan_file=stan)
         rdata = os.path.join(DATAFILES_PATH, 'eight_schools.data.R')
-        with LogCapture() as log:
+        with LogCapture(level=logging.WARNING) as log:
             logging.getLogger()
             fit = model.sample(
                 data=rdata,
                 seed=55157,
                 show_progress=False,
                 show_console=False,
             )
-
-            self.assertTrue(
-                any([a != b for a, b in zip(fit.max_treedepths, [0, 0, 0, 0])])
-            )
+            msg = log.actual()[-1][-1]
             self.assertTrue(
-                any([a != b for a, b in zip(fit.divergences, [0, 0, 0, 0])])
+                msg.startswith('Some chains may have failed to converge.')
             )
+            self.assertFalse(np.all(fit.divergences == 0))
 
-            log.check_present(
-                (
-                    'cmdstanpy',
-                    'WARNING',
-                    'Some chains may have failed to converge.\n'
-                    '\tChain 1 had 10 divergent transitions (1.0%)\n'
-                    '\tChain 2 had 143 divergent transitions (14.3%)\n'
-                    '\tChain 3 had 5 divergent transitions (0.5%)\n'
-                    '\tChain 4 had 4 divergent transitions (0.4%)\n'
-                    '\tChain 4 had 6 iterations at max treedepth (0.6%)\n'
-                    '\tUse function "diagnose()" to see further information.',
-                ),
+        with LogCapture(level=logging.WARNING) as log:
+            logging.getLogger()
+            fit = model.sample(
+                data=rdata,
+                seed=40508,
+                max_treedepth=3,
             )
+            msg = log.actual()[-1][-1]
+            self.assertTrue('max treedepth' in msg)
+            self.assertFalse(np.all(fit.max_treedepths == 0))
 
         stan = os.path.join(DATAFILES_PATH, 'bernoulli.stan')
         model = CmdStanModel(stan_file=stan)
         jdata = os.path.join(DATAFILES_PATH, 'bernoulli.data.json')
         fit = model.sample(
             data=jdata,
-            chains=2,
-            parallel_chains=2,
-            seed=12345,
             iter_warmup=200,
             iter_sampling=100,
-            show_progress=False,
-            show_console=False,
         )
         self.assertTrue(np.all(fit.divergences == 0))
         self.assertTrue(np.all(fit.max_treedepths == 0))