Repository: canonical/operator
Pull request #655 (2021-10-22 02:33:55)

Exec test fixes and "for line in process.stdout" fix
* Fix ResourceWarnings in TestRealPebble from not closing (web)sockets
* Avoid ResourceWarning in test_send_signal from not calling process.wait
* Suppress expected unhandled thread exception warning in two tests 
* Fix "for line in process.stdout" using TextIOWrapper 


Patch:
diff --git a/ops/pebble.py b/ops/pebble.py
index b5070f0..8c57d85 100644
--- a/ops/pebble.py
+++ b/ops/pebble.py
@@ -25,6 +25,7 @@ import enum
 import http.client
 import io
 import json
+import logging
 import os
 import re
 import select
@@ -43,6 +44,8 @@ import warnings
 from ops._private import yaml
 from ops._vendor import websocket
 
+logger = logging.getLogger(__name__)
+
 _not_provided = object()
 
 
@@ -806,6 +809,8 @@ class ExecProcess:
         client: 'Client',
         timeout: typing.Optional[float],
         control_ws: websocket.WebSocket,
+        stdio_ws: websocket.WebSocket,
+        stderr_ws: websocket.WebSocket,
         command: typing.List[str],
         encoding: typing.Optional[str],
         change_id: ChangeID,
@@ -819,6 +824,8 @@ class ExecProcess:
         self._client = client
         self._timeout = timeout
         self._control_ws = control_ws
+        self._stdio_ws = stdio_ws
+        self._stderr_ws = stderr_ws
         self._command = command
         self._encoding = encoding
         self._change_id = change_id
@@ -867,7 +874,11 @@ class ExecProcess:
         if self._cancel_reader is not None:
             os.close(self._cancel_reader)
 
-        self._control_ws.close()
+        # Close websockets (shutdown doesn't send CLOSE message or wait for response).
+        self._control_ws.shutdown()
+        self._stdio_ws.shutdown()
+        if self._stderr_ws is not None:
+            self._stderr_ws.shutdown()
 
         if change.err:
             raise ChangeError(change.err, change)
@@ -965,12 +976,15 @@ def _websocket_to_writer(ws, writer, encoding):
 
         if isinstance(chunk, str):
             try:
-                command = json.loads(chunk)
+                payload = json.loads(chunk)
             except ValueError:
                 # Garbage sent, try to keep going
+                logger.warning('Cannot decode I/O command (invalid JSON)')
                 continue
-            if command.get('command') != 'end':
+            command = payload.get('command')
+            if command != 'end':
                 # A command we don't recognize, keep going
+                logger.warning('Invalid I/O command {!r}'.format(command))
                 continue
             # Received "end" command (EOF signal), stop thread
             break
@@ -983,10 +997,8 @@ def _websocket_to_writer(ws, writer, encoding):
 class _WebsocketWriter(io.BufferedIOBase):
     """A writable file-like object that sends what's written to it to a websocket."""
 
-    def __init__(self, name, ws, encoding):
-        self.name = name
+    def __init__(self, ws):
         self.ws = ws
-        self.encoding = encoding
 
     def writable(self):
         """Denote this file-like object as writable."""
@@ -994,12 +1006,8 @@ class _WebsocketWriter(io.BufferedIOBase):
 
     def write(self, chunk):
         """Write chunk to the websocket."""
-        if isinstance(chunk, str):
-            if self.encoding is None:
-                raise ValueError('encoding must be set if writing str to {}'.format(self.name))
-            chunk = chunk.encode(self.encoding)
-        elif self.encoding is not None:
-            raise ValueError('encoding must be None if writing bytes to {}'.format(self.name))
+        if not isinstance(chunk, bytes):
+            raise TypeError('value to write must be bytes, not {}'.format(type(chunk).__name__))
         self.ws.send_binary(chunk)
         return len(chunk)
 
@@ -1011,10 +1019,10 @@ class _WebsocketWriter(io.BufferedIOBase):
 class _WebsocketReader(io.BufferedIOBase):
     """A readable file-like object whose reads come from a websocket."""
 
-    def __init__(self, ws, encoding):
+    def __init__(self, ws):
         self.ws = ws
-        self.encoding = encoding
-        self.remaining = '' if encoding is not None else b''
+        self.remaining = b''
+        self.eof = False
 
     def readable(self):
         """Denote this file-like object as readable."""
@@ -1022,17 +1030,30 @@ class _WebsocketReader(io.BufferedIOBase):
 
     def read(self, n=-1):
         """Read up to n bytes from the websocket (or one message if n<0)."""
-        if not self.remaining:
-            recv = self.ws.recv()
-
-            if isinstance(recv, str):
-                _ = json.loads(recv)  # raise ValueError on invalid JSON
+        if self.eof:
+            # Calling read() multiple times after EOF should still return EOF
+            return b''
+
+        while not self.remaining:
+            chunk = self.ws.recv()
+
+            if isinstance(chunk, str):
+                try:
+                    payload = json.loads(chunk)
+                except ValueError:
+                    # Garbage sent, try to keep going
+                    logger.warning('Cannot decode I/O command (invalid JSON)')
+                    continue
+                command = payload.get('command')
+                if command != 'end':
+                    # A command we don't recognize, keep going
+                    logger.warning('Invalid I/O command {!r}'.format(command))
+                    continue
                 # Received "end" command, return EOF designator
-                return '' if self.encoding is not None else b''
+                self.eof = True
+                return b''
 
-            if self.encoding is not None:
-                recv = recv.decode(self.encoding)
-            self.remaining = recv
+            self.remaining = chunk
 
         if n < 0:
             n = len(self.remaining)
@@ -1040,6 +1061,10 @@ class _WebsocketReader(io.BufferedIOBase):
         self.remaining = self.remaining[n:]
         return result
 
+    def read1(self, n=-1):
+        """An alias for read."""
+        return self.read(n)
+
 
 class Client:
     """Pebble API client."""
@@ -1847,14 +1872,18 @@ class Client:
             threads.append(t)
             process_stdin = None
         else:
-            process_stdin = _WebsocketWriter('stdin', stdio_ws, encoding)
+            process_stdin = _WebsocketWriter(stdio_ws)
+            if encoding is not None:
+                process_stdin = io.TextIOWrapper(process_stdin, encoding=encoding, newline='')
 
         if stdout is not None:
             t = _start_thread(_websocket_to_writer, stdio_ws, stdout, encoding)
             threads.append(t)
             process_stdout = None
         else:
-            process_stdout = _WebsocketReader(stdio_ws, encoding)
+            process_stdout = _WebsocketReader(stdio_ws)
+            if encoding is not None:
+                process_stdout = io.TextIOWrapper(process_stdout, encoding=encoding, newline='')
 
         process_stderr = None
         if not combine_stderr:
@@ -1862,7 +1891,10 @@ class Client:
                 t = _start_thread(_websocket_to_writer, stderr_ws, stderr, encoding)
                 threads.append(t)
             else:
-                process_stderr = _WebsocketReader(stderr_ws, encoding)
+                process_stderr = _WebsocketReader(stderr_ws)
+                if encoding is not None:
+                    process_stderr = io.TextIOWrapper(
+                        process_stderr, encoding=encoding, newline='')
 
         process = ExecProcess(
             stdin=process_stdin,
@@ -1870,6 +1902,8 @@ class Client:
             stderr=process_stderr,
             client=self,
             timeout=timeout,
+            stdio_ws=stdio_ws,
+            stderr_ws=stderr_ws,
             control_ws=control_ws,
             command=command,
             encoding=encoding,


Test patch:
diff --git a/test/test_pebble.py b/test/test_pebble.py
index a739230..54c8f49 100644
--- a/test/test_pebble.py
+++ b/test/test_pebble.py
@@ -29,6 +29,8 @@ import unittest
 import unittest.mock
 import unittest.util
 
+import pytest
+
 import ops.pebble as pebble
 from ops._private import yaml
 from ops._vendor import websocket
@@ -2081,7 +2083,7 @@ class MockWebsocket:
     def recv(self):
         return self.receives.pop(0)
 
-    def close(self):
+    def shutdown(self):
         pass
 
 
@@ -2248,8 +2250,11 @@ class TestExec(unittest.TestCase):
             process.send_signal(signal.SIGHUP)
             num_sends += 2
 
+        process.wait()
+
         self.assertEqual(self.client.requests, [
             ('POST', '/v1/exec', None, self.build_exec_data(['server'])),
+            ('GET', '/v1/changes/123/wait', {'timeout': '4.000s'}, None),
         ])
 
         self.assertEqual(len(control.sends), num_sends)
@@ -2391,6 +2396,31 @@ class TestExec(unittest.TestCase):
             ('TXT', '{"command":"end"}'),
         ])
 
+    def test_wait_output_bad_command(self):
+        stdio, stderr, _ = self.add_responses('123', 0)
+        stdio.receives.append(b'Python 3.8.10\n')
+        stdio.receives.append('not json')  # bad JSON should be ignored
+        stdio.receives.append('{"command":"foo"}')  # unknown command should be ignored
+        stdio.receives.append('{"command":"end"}')
+        stderr.receives.append('{"command":"end"}')
+
+        with self.assertLogs('ops.pebble', level='WARNING') as cm:
+            process = self.client.exec(['python3', '--version'])
+            out, err = process.wait_output()
+        self.assertEqual(cm.output, [
+            "WARNING:ops.pebble:Cannot decode I/O command (invalid JSON)",
+            "WARNING:ops.pebble:Invalid I/O command 'foo'",
+        ])
+
+        self.assertEqual(out, 'Python 3.8.10\n')
+        self.assertEqual(err, '')
+
+        self.assertEqual(self.client.requests, [
+            ('POST', '/v1/exec', None, self.build_exec_data(['python3', '--version'])),
+            ('GET', '/v1/changes/123/wait', {'timeout': '4.000s'}, None),
+        ])
+        self.assertEqual(stdio.sends, [])
+
     def test_wait_passed_output(self):
         io_ws, stderr, _ = self.add_responses('123', 0)
         io_ws.receives.append(b'foo\n')
@@ -2450,6 +2480,35 @@ class TestExec(unittest.TestCase):
         ])
         self.assertEqual(io_ws.sends, [])
 
+    def test_wait_passed_output_bad_command(self):
+        io_ws, stderr, _ = self.add_responses('123', 0)
+        io_ws.receives.append(b'foo\n')
+        io_ws.receives.append('not json')  # bad JSON should be ignored
+        io_ws.receives.append('{"command":"foo"}')  # unknown command should be ignored
+        io_ws.receives.append('{"command":"end"}')
+        stderr.receives.append(b'some error\n')
+        stderr.receives.append('{"command":"end"}')
+
+        out = io.StringIO()
+        err = io.StringIO()
+
+        with self.assertLogs('ops.pebble', level='WARNING') as cm:
+            process = self.client.exec(['echo', 'foo'], stdout=out, stderr=err)
+            process.wait()
+        self.assertEqual(cm.output, [
+            "WARNING:ops.pebble:Cannot decode I/O command (invalid JSON)",
+            "WARNING:ops.pebble:Invalid I/O command 'foo'",
+        ])
+
+        self.assertEqual(out.getvalue(), 'foo\n')
+        self.assertEqual(err.getvalue(), 'some error\n')
+
+        self.assertEqual(self.client.requests, [
+            ('POST', '/v1/exec', None, self.build_exec_data(['echo', 'foo'])),
+            ('GET', '/v1/changes/123/wait', {'timeout': '4.000s'}, None),
+        ])
+        self.assertEqual(io_ws.sends, [])
+
     @unittest.skipIf(sys.platform == 'win32', "exec() with files doesn't work on Windows")
     def test_wait_file_io(self):
         fin = tempfile.TemporaryFile(mode='w+', encoding='utf-8')
@@ -2495,9 +2554,8 @@ class TestExec(unittest.TestCase):
         process = self.client.exec(['awk', '{ print toupper($) }'])
         process.stdin.write('Foo Bar\n')
         self.assertEqual(process.stdout.read(4), 'FOO ')
-        self.assertEqual(process.stdout.read(), 'BAR\n')
         process.stdin.write('bazz\n')
-        self.assertEqual(process.stdout.read(), 'BAZZ\n')
+        self.assertEqual(process.stdout.read(), 'BAR\nBAZZ\n')
         process.stdin.close()
         self.assertEqual(process.stdout.read(), '')
         process.wait()
@@ -2507,8 +2565,7 @@ class TestExec(unittest.TestCase):
             ('GET', '/v1/changes/123/wait', {'timeout': '4.000s'}, None),
         ])
         self.assertEqual(stdio.sends, [
-            ('BIN', b'Foo Bar\n'),
-            ('BIN', b'bazz\n'),
+            ('BIN', b'Foo Bar\nbazz\n'),  # TextIOWrapper groups the writes together
             ('TXT', '{"command":"end"}'),
         ])
 
@@ -2555,7 +2612,7 @@ class TestExec(unittest.TestCase):
             self.client.exec(['foo'])
         self.assertIn(str(cm.exception), 'unexpected error connecting to websockets: conn!')
 
-    def test_websocket_send_binary_raises(self):
+    def test_websocket_send_raises(self):
         stdio, stderr, _ = self.add_responses('123', 0)
         raised = False
 
@@ -2580,6 +2637,12 @@ class TestExec(unittest.TestCase):
         ])
         self.assertEqual(stdio.sends, [])
 
+    # You'd normally use pytest.mark.filterwarnings as a decorator, but
+    # PytestUnhandledThreadExceptionWarning isn't present on older Python versions.
+    if hasattr(pytest, 'PytestUnhandledThreadExceptionWarning'):
+        test_websocket_send_raises = pytest.mark.filterwarnings(
+            'ignore::pytest.PytestUnhandledThreadExceptionWarning')(test_websocket_send_raises)
+
     def test_websocket_recv_raises(self):
         stdio, stderr, _ = self.add_responses('123', 0)
         raised = False
@@ -2607,6 +2670,10 @@ class TestExec(unittest.TestCase):
             ('TXT', '{"command":"end"}'),
         ])
 
+    if hasattr(pytest, 'PytestUnhandledThreadExceptionWarning'):
+        test_websocket_recv_raises = pytest.mark.filterwarnings(
+            'ignore::pytest.PytestUnhandledThreadExceptionWarning')(test_websocket_recv_raises)
+
 
 # Set the RUN_REAL_PEBBLE_TESTS environment variable to run these tests
 # against a real Pebble server. For example, in one terminal, run Pebble:
@@ -2693,23 +2760,42 @@ class TestRealPebble(unittest.TestCase):
         process = self.client.exec(['cat'])
 
         def stdin_thread():
-            for line in ['one\n', '2\n', 'THREE\n']:
-                process.stdin.write(line)
-                process.stdin.flush()
-                time.sleep(0.1)
-            process.stdin.close()
+            try:
+                for line in ['one\n', '2\n', 'THREE\n']:
+                    process.stdin.write(line)
+                    process.stdin.flush()
+                    time.sleep(0.1)
+            finally:
+                process.stdin.close()
 
         threading.Thread(target=stdin_thread).start()
 
-        # TODO: fix OSError with "for line in process.stdout: ...":
-        # OSError: read() should have returned a bytes object, not 'str'
         reads = []
-        while True:
-            chunk = process.stdout.read()
-            if not chunk:
-                break
-            reads.append(chunk)
+        for line in process.stdout:
+            reads.append(line)
 
         process.wait()
 
         self.assertEqual(reads, ['one\n', '2\n', 'THREE\n'])
+
+    def test_exec_streaming_bytes(self):
+        process = self.client.exec(['cat'], encoding=None)
+
+        def stdin_thread():
+            try:
+                for line in [b'one\n', b'2\n', b'THREE\n']:
+                    process.stdin.write(line)
+                    process.stdin.flush()
+                    time.sleep(0.1)
+            finally:
+                process.stdin.close()
+
+        threading.Thread(target=stdin_thread).start()
+
+        reads = []
+        for line in process.stdout:
+            reads.append(line)
+
+        process.wait()
+
+        self.assertEqual(reads, [b'one\n', b'2\n', b'THREE\n'])