Create 2 unit tests for wait_for_generate

diff --git a/test/test_arnied_wrapper.py b/test/test_arnied_wrapper.py
index 2e2476e..3d5265e 100755 (executable)
--- a/test/test_arnied_wrapper.py
+++ b/test/test_arnied_wrapper.py
@@ -22,7 +22,7 @@
 
 import unittest
 import unittest.mock as mock
-import subprocess
+import time
 
 from src import arnied_wrapper
 
@@ -107,3 +107,23 @@ class ArniedWrapperTest(unittest.TestCase):
         """Test e-mail transferring."""
         DummyCmdOutputMapping.asserted_cmds = [{"cmd": '/usr/intranator/bin/arnied_helper --transfer-mail', "stdout": b"", "returncode": 0}]
         arnied_wrapper.email_transfer()
+
+    def test_wait_for_no_generate(self):
+        """Test waiting for generate if there is no generate scheduled"""
+        DummyCmdOutputMapping.asserted_cmds = [
+            {"cmd": f'/usr/intranator/bin/arnied_helper --is-scheduled-or-running {job}', "stdout": b"", "returncode": 1}
+                for job in ("GENERATE", "GENERATE_OFFLINE")]
+        start_time = time.time()
+        self.assertEqual(arnied_wrapper.wait_for_generate(),  True)
+        self.assertLess(time.time() - start_time, 1)    # no sleep, just call and return at once
+
+    def test_wait_for_generate(self):
+        """Test waiting for generate if there is generate scheduled"""
+        DummyCmdOutputMapping.asserted_cmds = [
+            {"cmd": f'/usr/intranator/bin/arnied_helper --is-scheduled-or-running {job}', "stdout": b"", "returncode": 0}
+            for job in ("GENERATE", "GENERATE_OFFLINE")]
+        start_time = time.time()
+        self.assertEqual(arnied_wrapper.wait_for_generate(timeout=2),  False)
+        end_time = time.time()
+        self.assertGreater(end_time - start_time, 1.5)
+        self.assertLess(end_time - start_time, 2.5)