Increase coverage

Author: qianduoduo0904

mars/oscar/backends/communication/socket.py

@@ -159,6 +159,7 @@ async def stop(self):
         await asyncio.gather(
             *(channel.close() for channel in self._channels if not channel.closed)
         )
+        self._channels = []
 
     @property
     @implements(Server.stopped)

mars/oscar/backends/communication/ucx.py

@@ -308,7 +308,7 @@ async def recv_buffers(self, buffers: list):
             try:
                 for buffer in buffers:
                     await self.ucp_endpoint.recv(buffer)
-            except BaseException as e:
+            except BaseException as e:  # pragma: no cover
                 if not self._closed:
                     # In addition to UCX exceptions, may be CancelledError or another
                     # "low-level" exception. The only safe thing to do is to abort.
@@ -455,6 +455,7 @@ async def stop(self):
         await asyncio.gather(
             *(channel.close() for channel in self._channels if not channel.closed)
         )
+        self._channels = []
         self._ucp_listener = None
         self._closed.set()
 

mars/oscar/backends/core.py

@@ -39,7 +39,7 @@ def __init__(self):
         self._clients: Dict[Client, asyncio.Task] = dict()
 
     async def get_client(self, router: Router, dest_address: str) -> Client:
-        client = await router.get_client(dest_address, from_who=self)
+        client = await router.get_client(dest_address, from_who=type(self))
         if client not in self._clients:
             self._clients[client] = asyncio.create_task(self._listen(client))
             self._client_to_message_futures[client] = dict()

mars/oscar/backends/router.py

@@ -191,7 +191,7 @@ async def get_client_via_type(
                     return cached_client
 
         client_type_to_addresses = self._get_client_type_to_addresses(external_address)
-        if client_type not in client_type_to_addresses:
+        if client_type not in client_type_to_addresses:  # pragma: no cover
             raise ValueError(
                 f"Client type({client_type}) is not supported for {external_address}"
             )

mars/oscar/backends/transfer.py

@@ -73,11 +73,11 @@ def __init__(
     @staticmethod
     @abstractmethod
     def _new_message(ref: Union[BufferRef, FileObjectRef], buf: Any):
-        pass
+        """new message"""
 
     @abstractmethod
     async def _read(self, index: int, buffer_or_fileobj: Any):
-        pass
+        """read data"""
 
     async def _send_one(
         self,
@@ -142,14 +142,14 @@ def __init__(
 
     @abstractmethod
     async def _write(self, index: int, buffer_or_fileobj: Any):
-        pass
+        """write data"""
 
     @staticmethod
     @staticmethod
     def _get_ref_from_message(
         message: Union[CopytoBuffersMessage, CopytoFileObjectsMessage]
     ):
-        pass
+        """get ref according to message"""
 
     async def _recv_part(self, buf: Any, index: int, message_id: bytes):
         async with _catch_error(self.channel, message_id):
@@ -285,7 +285,7 @@ async def copyto_via_buffers(
             # do not support buffer copy
             # send data in batches
             client, is_cached = await router.get_client(
-                address, from_who=self, return_from_cache=True
+                address, from_who=type(self), return_from_cache=True
             )
             if not is_cached:
                 # tell server to switch to transfer dedicated channel
@@ -299,7 +299,7 @@ async def copyto_via_buffers(
                 )
         else:
             client, is_cached = await router.get_client_via_type(
-                address, client_type, from_who=self, return_from_cache=True
+                address, client_type, from_who=type(self), return_from_cache=True
             )
             if not is_cached:
                 # tell server to switch to transfer dedicated channel
@@ -339,7 +339,7 @@ async def copyto_via_file_objects(
         ), "`copyto_via_file_objects` can only be used inside pools"
         address = remote_file_object_refs[0].address
         client, is_cached = await router.get_client(
-            address, from_who=self, return_from_cache=True
+            address, from_who=type(self), return_from_cache=True
         )
         if not is_cached:
             # tell server to switch to transfer dedicated channel