traffic_learner: return an error code

And use it in tests, rather than expecting exact strings.

Signed-off-by: Douglas Bagnall <douglas.bagnall@catalyst.net.nz>
Reviewed-by: Andrew Bartlett <abartlet@samba.org>

diff --git a/python/samba/tests/blackbox/traffic_learner.py b/python/samba/tests/blackbox/traffic_learner.py
index 3fc35324f0f018778662be1bb62b6b3b7e4c00af..dd7c7c15678824a7bfe6801caa7b54152acebc2d 100644 (file)
--- a/python/samba/tests/blackbox/traffic_learner.py
+++ b/python/samba/tests/blackbox/traffic_learner.py
@@ -43,11 +43,8 @@ def temp_file(temp_dir):
 class TrafficLearnerTests(BlackboxTestCase):
 
     def test_no_output_file(self):
-        """Run the script with no output file specified"""
-        expected = (b"No output file was specified to write the model to.\n"
-                    b"Please specify a filename using the --out option.\n")
-        actual = self.check_output(LEARNER)
-        self.assertEquals(expected, actual)
+        """Run the script with no output file specified. Should fail."""
+        self.check_exit_code(LEARNER, 1)
 
     def test_model_generation(self):
         """Ensure a model is generated from a summary file and it is

diff --git a/script/traffic_learner b/script/traffic_learner
index 22341ef231c33ad0f3029813023c40170245c097..b3570d45eeab3bb4d2870a6208a991097420fff8 100755 (executable)
--- a/script/traffic_learner
+++ b/script/traffic_learner
@@ -36,10 +36,10 @@ def main():
                         help="read from this file (default STDIN)")
     args = parser.parse_args()
 
-    if not args.out:
+    if args.out is None:
         print("No output file was specified to write the model to.", file=sys.stdout)
         print("Please specify a filename using the --out option.", file=sys.stdout)
-        return
+        return 1
 
     if args.SUMMARY_FILE is sys.stdin:
         print("reading from STDIN...", file=sys.stderr)
@@ -59,5 +59,4 @@ def main():
 
     model.save(args.out)
 
-
-main()
+sys.exit(main())