From a84d0e27912ebd3e71947a456a7dafafd31b8b38 Mon Sep 17 00:00:00 2001
From: Juan Jose Nicola <juan.nicola@greenbone.net>
Date: Fri, 9 Oct 2020 10:54:48 +0200
Subject: [PATCH] Add test

---
 tests/test_scan_and_result.py | 40 ++++++++++++++++++++++++++++++-----
 1 file changed, 35 insertions(+), 5 deletions(-)

diff --git a/tests/test_scan_and_result.py b/tests/test_scan_and_result.py
index 9ab57c71..f53e6c01 100644
--- a/tests/test_scan_and_result.py
+++ b/tests/test_scan_and_result.py
@@ -238,7 +238,8 @@ def test_get_vts_filter_negative(self):
         )
         fs = FakeStream()
         self.daemon.handle_command(
-            '<get_vts filter="modification_time&lt;19000203"></get_vts>', fs,
+            '<get_vts filter="modification_time&lt;19000203"></get_vts>',
+            fs,
         )
         response = fs.get_response()
 
@@ -741,7 +742,8 @@ def test_get_scan_results_clean(self):
 
         fs = FakeStream()
         self.daemon.handle_command(
-            '<get_scans scan_id="%s" pop_results="1"/>' % scan_id, fs,
+            '<get_scans scan_id="%s" pop_results="1"/>' % scan_id,
+            fs,
         )
 
         res_len = len(
@@ -771,7 +773,8 @@ def test_get_scan_results_restore(self):
 
         fs = FakeStream(return_value=False)
         self.daemon.handle_command(
-            '<get_scans scan_id="%s" pop_results="1"/>' % scan_id, fs,
+            '<get_scans scan_id="%s" pop_results="1"/>' % scan_id,
+            fs,
         )
 
         res_len = len(
@@ -1060,6 +1063,31 @@ def test_get_scan_without_scanid(self):
             fs,
         )
 
+    def test_set_scan_total_hosts(self):
+
+        fs = FakeStream()
+        self.daemon.handle_command(
+            '<start_scan parallel="2">'
+            '<scanner_params />'
+            '<targets><target>'
+            '<hosts>localhost1, localhost2, localhost3, localhost4</hosts>'
+            '<ports>22</ports>'
+            '</target></targets>'
+            '</start_scan>',
+            fs,
+        )
+        self.daemon.start_queued_scans()
+
+        response = fs.get_response()
+        scan_id = response.findtext('id')
+
+        count = self.daemon.scan_collection.get_count_total(scan_id)
+        self.assertEqual(count, 4)
+
+        self.daemon.set_scan_total_hosts(scan_id, 3)
+        count = self.daemon.scan_collection.get_count_total(scan_id)
+        self.assertEqual(count, 3)
+
     def test_get_scan_progress_xml(self):
 
         fs = FakeStream()
@@ -1087,7 +1115,8 @@ def test_get_scan_progress_xml(self):
 
         fs = FakeStream()
         self.daemon.handle_command(
-            '<get_scans scan_id="%s" details="0" progress="1"/>' % scan_id, fs,
+            '<get_scans scan_id="%s" details="0" progress="1"/>' % scan_id,
+            fs,
         )
         response = fs.get_response()
 
@@ -1164,7 +1193,8 @@ def test_scan_exists(self, mock_create_process, _mock_os):
         )
 
         self.daemon.handle_command(
-            cmd, fs,
+            cmd,
+            fs,
         )
         self.daemon.start_queued_scans()