python/samba/tests/blackbox/traffic_learner.py

   1 # Black box tests for script/traffic_leaner
   2 #
   3 # Copyright (C) Catalyst IT Ltd. 2017
   4 #
   5 # This program is free software; you can redistribute it and/or modify
   6 # it under the terms of the GNU General Public License as published by
   7 # the Free Software Foundation; either version 3 of the License, or
   8 # (at your option) any later version.
   9 #
  10 # This program is distributed in the hope that it will be useful,
  11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13 # GNU General Public License for more details.
  14 #
  15 # You should have received a copy of the GNU General Public License
  16 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  17
  18 """Blackbox tests for traffic_leaner"""
  19
  20 from contextlib import contextmanager
  21 import os
  22 import tempfile
  23
  24 from samba.tests import BlackboxTestCase
  25
  26 LEARNER  = "script/traffic_learner"
  27 DATA_DIR = "python/samba/tests/blackbox/testdata"
  28
  29
  30 @contextmanager
  31 def temp_file(temp_dir):
  32     try:
  33         tf   = tempfile.NamedTemporaryFile(dir=temp_dir)
  34         name = tf.name
  35         tf.close()
  36         yield name
  37     finally:
  38         if os.path.exists(name):
  39             os.remove(name)
  40
  41
  42 class TrafficLearnerTests(BlackboxTestCase):
  43
  44     def test_no_output_file(self):
  45         """Run the script with no output file specified"""
  46         expected = ("No output file was specified to write the model to.\n"
  47                     "Please specify a filename using the --out option.\n")
  48         actual = self.check_output(LEARNER)
  49         self.assertEquals(expected, actual)
  50
  51     def test_model_generation(self):
  52         """Ensure a model is generated from a summary file and it is
  53            correct"""
  54
  55         with temp_file(self.tempdir) as output:
  56             summary  = os.path.join(DATA_DIR, "traffic-sample-very-short.txt")
  57             command  = "%s %s --out %s" % (LEARNER, summary, output)
  58             self.check_run(command)
  59             expected_fn = os.path.join(DATA_DIR, "traffic_learner.expected")
  60             expected = open(expected_fn).readlines()
  61             actual = open(output).readlines()
  62             self.assertEquals(expected, actual)