Added session.print (#161)

* added print to the console
* fixed task leak

molotov/_grpc.py

@@ -8,4 +8,12 @@ def grpc_session(loop, console, verbose, statsd, trace_config, **kw):
     url = kw["grpc_url"]
     channel = aio.insecure_channel(url)
     channel._trace_configs = [trace_config]  # type: ignore
+
+    def _print(data):
+        if console is None:
+            print(data)
+        else:
+            console.print(data)
+
+    channel.print = _print  # type: ignore
     return channel

molotov/session.py

@@ -103,13 +103,24 @@ def get_session(loop, console, verbose=0, statsd=None, kind="http", **kw):
     request_class.verbose = verbose
     request_class.response_class = LoggedClientResponse
     request_class.tracer = trace_config
+
+    # patching the class to avoid aiohttp warning
+    def _print(self, data):
+        if self.console is None:
+            print(data)
+        else:
+            self.console.print(data)
+
+    ClientSession.print = _print  # type: ignore
+
     session = ClientSession(
         request_class=request_class,
         response_class=LoggedClientResponse,
         connector=connector,
         trace_configs=[trace_config],
         **kw,
     )
+    session.console = console
 
     return session
 

molotov/tests/example10.py

@@ -8,7 +8,10 @@
 
 
 @scenario(weight=40)
-async def grpc_scenario(session, session_factory="grpc", grpc_url="ipv4:///127.0.0.1:50051"):
+async def grpc_scenario(
+    session, session_factory="grpc", grpc_url="ipv4:///127.0.0.1:50051"
+):
     stub = helloworld_pb2_grpc.GreeterStub(session)
     response = await stub.SayHello(helloworld_pb2.HelloRequest(name="Alice"))
+    session.print("Call done")
     assert response.message == "Hello, Alice!", response.message

molotov/tests/example11.py

@@ -0,0 +1,6 @@
+import molotov
+
+
+@molotov.scenario()
+async def test_print(session):
+    session.print("Hello")

molotov/tests/support.py

@@ -171,7 +171,9 @@ async def read(self):
 
         def feed_data(self, data):
             if body == b"":
-                err = AttributeError("'EmptyStreamReader' object has no " "attribute 'unread_data'")
+                err = AttributeError(
+                    "'EmptyStreamReader' object has no " "attribute 'unread_data'"
+                )
                 raise err
             pass
 
@@ -363,7 +365,9 @@ async def _slept(delay, result=None, *, loop=None):
         # forces a context switch
         await original(0)
 
-    with patch("asyncio.sleep", _slept):
+    with patch("asyncio.sleep", _slept), patch(
+        "molotov.util.cancellable_sleep", _slept
+    ):
         yield calls
 
 

molotov/tests/test_run.py

@@ -454,7 +454,9 @@ async def sizer(session):
             stdout, stderr = stdout.read().strip(), stderr.read().strip()
 
             # stdout, stderr, rc = self._test_molotov()
-            ratio = float(counters["FAILED"].value) / float(counters["OK"].value) * 100.0
+            ratio = (
+                float(counters["FAILED"].value) / float(counters["OK"].value) * 100.0
+            )
             self.assertTrue(ratio >= 4.75, ratio)
 
     @co_catch_output
@@ -805,7 +807,9 @@ async def scenario_one(session):
                 }
                 data = json.dumps({"1": "xxx"})
 
-                with aiohttp.MultipartWriter("form-data", boundary=boundary) as mpwriter:
+                with aiohttp.MultipartWriter(
+                    "form-data", boundary=boundary
+                ) as mpwriter:
                     mpwriter.append(
                         data,
                         {
@@ -847,4 +851,24 @@ def test_grpc_factory(self):
 
         with grpc_server():
             stdout, stderr, rc = self._test_molotov("--max-runs", "1", test)
+
+        self.assertTrue("Call done" in stdout, stdout)
         self.assertTrue("SUCCESSES: 1" in stdout, stdout)
+
+    @co_catch_output
+    @dedicatedloop
+    def test_print_ui(self):
+        test = os.path.join(_HERE, "example11.py")
+
+        with coserver():
+            stdout, stderr, rc = self._test_molotov("--max-runs", "1", test)
+        self.assertTrue("Hello" in stdout, stdout)
+
+    @co_catch_output
+    @dedicatedloop
+    def test_print_console(self):
+        test = os.path.join(_HERE, "example11.py")
+
+        with coserver():
+            stdout, stderr, rc = self._test_molotov("--max-runs", "1", "-c", test)
+        self.assertTrue("Hello" in stdout, stdout)

molotov/ui/app.py

@@ -20,8 +20,11 @@
     TerminalController,
     create_key_bindings,
 )
+from molotov.util import cancellable_sleep
 
-TITLE = HTML(f"<b>Molotov v{__version__}</b> ~ Happy Breaking 🥛🔨 ~ <i>Ctrl+C to abort</i>")
+TITLE = HTML(
+    f"<b>Molotov v{__version__}</b> ~ Happy Breaking 🥛🔨 ~ <i>Ctrl+C to abort</i>"
+)
 
 
 class MolotovApp:
@@ -61,19 +64,22 @@ def __init__(
         self.refresh_interval = refresh_interval
         self._running = False
 
+    def _dump_term(self):
+        for line in self.terminal.dump():
+            sys.stdout.write(line)
+            sys.stdout.flush()
+
+    def _dump_errors(self):
+        for line in self.errors.dump():
+            sys.stdout.write(line)
+            sys.stdout.flush()
+
     async def refresh_console(self):
         while self._running:
-            for line in self.terminal.dump():
-                sys.stdout.write(line)
-                sys.stdout.flush()
-
+            self._dump_term()
             await asyncio.sleep(0)
-
-            for line in self.errors.dump():
-                sys.stdout.write(line)
-                sys.stdout.flush()
-
-            await asyncio.sleep(self.refresh_interval)
+            self._dump_errors()
+            await cancellable_sleep(self.refresh_interval)
 
     async def start(self):
         self._running = True
@@ -126,7 +132,11 @@ def _handle_exception(*args, **kw):
 
     async def stop(self):
         self._running = False
-        await asyncio.sleep(0)
+        cancellable_sleep.cancel_all()
+
+        # dump any remaining data
+        self._dump_term()
+        self._dump_errors()
 
         if not self.simple_console:
             try:

molotov/worker.py

@@ -184,6 +184,7 @@ async def _run(self):
                 self.print(f"Ran {self.count} scenarios")
             step_start = now()
             result = await self.step(self.count, scenario=single, options=options)
+
             if result == 1:
                 self.results["OK"] += 1
                 self.results["MINUTE_OK"] += 1