Initial commit: Core packages

odoo-bringout-oca-ocb-bus/bus/tests/common.py

@@ -0,0 +1,124 @@
+# Part of Odoo. See LICENSE file for full copyright and licensing details.
+
+import json
+import struct
+from threading import Event
+import unittest
+from unittest.mock import patch
+
+try:
+    import websocket
+except ImportError:
+    websocket = None
+
+import odoo.tools
+from odoo.tests import HOST, HttpCase, TEST_CURSOR_COOKIE_NAME
+from ..websocket import CloseCode, Websocket, WebsocketConnectionHandler
+
+
+class WebsocketCase(HttpCase):
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        if websocket is None:
+            cls._logger.warning("websocket-client module is not installed")
+            raise unittest.SkipTest("websocket-client module is not installed")
+        cls._WEBSOCKET_URL = f"ws://{HOST}:{odoo.tools.config['http_port']}/websocket"
+        websocket_allowed_patch = patch.object(WebsocketConnectionHandler, "websocket_allowed", return_value=True)
+        cls.startClassPatcher(websocket_allowed_patch)
+
+    def setUp(self):
+        super().setUp()
+        self._websockets = set()
+        # Used to ensure websocket connections have been closed
+        # properly.
+        self._websocket_events = set()
+        original_serve_forever = WebsocketConnectionHandler._serve_forever
+
+        def _mocked_serve_forever(*args):
+            websocket_closed_event = Event()
+            self._websocket_events.add(websocket_closed_event)
+            original_serve_forever(*args)
+            websocket_closed_event.set()
+
+        self._serve_forever_patch = patch.object(
+            WebsocketConnectionHandler,
+            '_serve_forever',
+            wraps=_mocked_serve_forever
+        )
+        self.startPatcher(self._serve_forever_patch)
+
+    def tearDown(self):
+        self._close_websockets()
+        super().tearDown()
+
+    def _close_websockets(self):
+        """
+        Close all the connected websockets and wait for the connection
+        to terminate.
+        """
+        for ws in self._websockets:
+            if ws.connected:
+                ws.close(CloseCode.CLEAN)
+        self.wait_remaining_websocket_connections()
+
+    def websocket_connect(self, *args, **kwargs):
+        """
+        Connect a websocket. If no cookie is given, the connection is
+        opened with a default session. The created websocket is closed
+        at the end of the test.
+        """
+        if 'cookie' not in kwargs:
+            self.session = self.authenticate(None, None)
+            kwargs['cookie'] = f'session_id={self.session.sid}'
+        kwargs['cookie'] += f';{TEST_CURSOR_COOKIE_NAME}={self.http_request_key}'
+        kwargs['timeout'] = 10  # keep a large timeout to avoid aving a websocket request escaping the test
+        ws = websocket.create_connection(
+            type(self)._WEBSOCKET_URL, *args, **kwargs
+        )
+        ws.ping()
+        ws.recv_data_frame(control_frame=True) # pong
+        self._websockets.add(ws)
+        return ws
+
+    def subscribe(self, websocket, channels=None, last=None, wait_for_dispatch=True):
+        """ Subscribe the websocket to the given channels.
+        :param websocket: The websocket of the client.
+        :param channels: The list of channels to subscribe to.
+        :param last: The last notification id the client received.
+        :param wait_for_dispatch: Whether to wait for the notification
+            dispatching trigerred by the subscription.
+        """
+        dispatch_bus_notification_done = Event()
+        original_dispatch_bus_notifications = Websocket._dispatch_bus_notifications
+
+        def _mocked_dispatch_bus_notifications(self, *args):
+            original_dispatch_bus_notifications(self, *args)
+            dispatch_bus_notification_done.set()
+
+        with patch.object(Websocket, '_dispatch_bus_notifications', _mocked_dispatch_bus_notifications):
+            sub = {'event_name': 'subscribe', 'data': {
+                'channels': channels or [],
+            }}
+            if last:
+                sub['data']['last'] = last
+            websocket.send(json.dumps(sub))
+            if wait_for_dispatch:
+                dispatch_bus_notification_done.wait(timeout=5)
+
+    def wait_remaining_websocket_connections(self):
+        """ Wait for the websocket connections to terminate. """
+        for event in self._websocket_events:
+            event.wait(5)
+
+    def assert_close_with_code(self, websocket, expected_code):
+        """
+        Assert that the websocket is closed with the expected_code.
+        """
+        opcode, payload = websocket.recv_data()
+        # ensure it's a close frame
+        self.assertEqual(opcode, 8)
+        code = struct.unpack('!H', payload[:2])[0]
+        # ensure the close code is the one we expected
+        self.assertEqual(code, expected_code)