pyetp 0.0.45__py3-none-any.whl → 0.0.47__py3-none-any.whl

pyetp/client.py

@@ -2,12 +2,11 @@ import asyncio
 import contextlib
 import datetime
 import logging
-import sys
 import typing as T
 import uuid
 import warnings
 from collections import defaultdict
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Generator
 from types import TracebackType
 
 import numpy as np
@@ -17,177 +16,94 @@ import websockets.client
 from etpproto.connection import CommunicationProtocol, ConnectionType, ETPConnection
 from etpproto.messages import Message, MessageFlags
 from etptypes import ETPModel
-from etptypes.energistics.etp.v12.datatypes.any_array_type import AnyArrayType
-from etptypes.energistics.etp.v12.datatypes.any_logical_array_type import (
+from pydantic import SecretStr
+from xtgeo import RegularSurface
+
+import resqml_objects.v201 as ro
+from energistics.etp.v12.datatypes import (
+    AnyArrayType,
     AnyLogicalArrayType,
+    ArrayOfString,
+    DataValue,
+    ErrorInfo,
+    SupportedDataObject,
+    SupportedProtocol,
+    Uuid,
+    Version,
 )
-from etptypes.energistics.etp.v12.datatypes.array_of_string import ArrayOfString
-from etptypes.energistics.etp.v12.datatypes.data_array_types.data_array_identifier import (
+from energistics.etp.v12.datatypes.data_array_types import (
     DataArrayIdentifier,
-)
-from etptypes.energistics.etp.v12.datatypes.data_array_types.data_array_metadata import (
     DataArrayMetadata,
-)
-from etptypes.energistics.etp.v12.datatypes.data_array_types.get_data_subarrays_type import (
     GetDataSubarraysType,
-)
-from etptypes.energistics.etp.v12.datatypes.data_array_types.put_data_arrays_type import (
     PutDataArraysType,
-)
-from etptypes.energistics.etp.v12.datatypes.data_array_types.put_data_subarrays_type import (
     PutDataSubarraysType,
-)
-from etptypes.energistics.etp.v12.datatypes.data_array_types.put_uninitialized_data_array_type import (
     PutUninitializedDataArrayType,
 )
-from etptypes.energistics.etp.v12.datatypes.data_value import DataValue
-from etptypes.energistics.etp.v12.datatypes.error_info import ErrorInfo
-from etptypes.energistics.etp.v12.datatypes.object.context_info import ContextInfo
-from etptypes.energistics.etp.v12.datatypes.object.context_scope_kind import (
+from energistics.etp.v12.datatypes.object import (
+    ContextInfo,
     ContextScopeKind,
-)
-from etptypes.energistics.etp.v12.datatypes.object.data_object import DataObject
-from etptypes.energistics.etp.v12.datatypes.object.dataspace import Dataspace
-from etptypes.energistics.etp.v12.datatypes.object.relationship_kind import (
+    DataObject,
+    Dataspace,
     RelationshipKind,
+    Resource,
 )
-from etptypes.energistics.etp.v12.datatypes.object.resource import Resource
-from etptypes.energistics.etp.v12.datatypes.supported_data_object import (
-    SupportedDataObject,
-)
-from etptypes.energistics.etp.v12.datatypes.supported_protocol import SupportedProtocol
-from etptypes.energistics.etp.v12.datatypes.uuid import Uuid
-from etptypes.energistics.etp.v12.datatypes.version import Version
-from etptypes.energistics.etp.v12.protocol.core.authorize import Authorize
-from etptypes.energistics.etp.v12.protocol.core.authorize_response import (
+from energistics.etp.v12.protocol.core import (
+    Authorize,
     AuthorizeResponse,
-)
-from etptypes.energistics.etp.v12.protocol.core.close_session import CloseSession
-from etptypes.energistics.etp.v12.protocol.core.open_session import OpenSession
-from etptypes.energistics.etp.v12.protocol.core.protocol_exception import (
+    CloseSession,
+    OpenSession,
     ProtocolException,
+    RequestSession,
 )
-from etptypes.energistics.etp.v12.protocol.core.request_session import RequestSession
-from etptypes.energistics.etp.v12.protocol.data_array.get_data_array_metadata import (
+from energistics.etp.v12.protocol.data_array import (
     GetDataArrayMetadata,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.get_data_array_metadata_response import (
     GetDataArrayMetadataResponse,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.get_data_arrays import (
     GetDataArrays,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.get_data_arrays_response import (
     GetDataArraysResponse,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.get_data_subarrays import (
     GetDataSubarrays,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.get_data_subarrays_response import (
     GetDataSubarraysResponse,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.put_data_arrays import (
     PutDataArrays,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.put_data_arrays_response import (
     PutDataArraysResponse,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.put_data_subarrays import (
     PutDataSubarrays,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.put_data_subarrays_response import (
     PutDataSubarraysResponse,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.put_uninitialized_data_arrays import (
     PutUninitializedDataArrays,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.put_uninitialized_data_arrays_response import (
     PutUninitializedDataArraysResponse,
 )
-from etptypes.energistics.etp.v12.protocol.dataspace.delete_dataspaces import (
+from energistics.etp.v12.protocol.dataspace import (
     DeleteDataspaces,
-)
-from etptypes.energistics.etp.v12.protocol.dataspace.delete_dataspaces_response import (
     DeleteDataspacesResponse,
-)
-from etptypes.energistics.etp.v12.protocol.dataspace.get_dataspaces import GetDataspaces
-from etptypes.energistics.etp.v12.protocol.dataspace.get_dataspaces_response import (
+    GetDataspaces,
     GetDataspacesResponse,
-)
-from etptypes.energistics.etp.v12.protocol.dataspace.put_dataspaces import PutDataspaces
-from etptypes.energistics.etp.v12.protocol.dataspace.put_dataspaces_response import (
+    PutDataspaces,
     PutDataspacesResponse,
 )
-from etptypes.energistics.etp.v12.protocol.discovery.get_resources import GetResources
-from etptypes.energistics.etp.v12.protocol.store.delete_data_objects import (
-    DeleteDataObjects,
+from energistics.etp.v12.protocol.discovery import (
+    GetResources,
+    GetResourcesResponse,
 )
-from etptypes.energistics.etp.v12.protocol.store.delete_data_objects_response import (
+from energistics.etp.v12.protocol.store import (
+    DeleteDataObjects,
     DeleteDataObjectsResponse,
-)
-from etptypes.energistics.etp.v12.protocol.store.get_data_objects import GetDataObjects
-from etptypes.energistics.etp.v12.protocol.store.get_data_objects_response import (
+    GetDataObjects,
     GetDataObjectsResponse,
-)
-from etptypes.energistics.etp.v12.protocol.store.put_data_objects import PutDataObjects
-from etptypes.energistics.etp.v12.protocol.store.put_data_objects_response import (
+    PutDataObjects,
     PutDataObjectsResponse,
 )
-from etptypes.energistics.etp.v12.protocol.transaction.commit_transaction import (
+from energistics.etp.v12.protocol.transaction import (
     CommitTransaction,
-)
-from etptypes.energistics.etp.v12.protocol.transaction.rollback_transaction import (
     RollbackTransaction,
-)
-from etptypes.energistics.etp.v12.protocol.transaction.start_transaction import (
     StartTransaction,
+    StartTransactionResponse,
 )
-from pydantic import SecretStr
-from xtgeo import RegularSurface
-
-import resqml_objects.v201 as ro
 from pyetp import utils_arrays, utils_xml
 from pyetp._version import version
 from pyetp.config import SETTINGS
+from pyetp.errors import ETPTransactionFailure
 from pyetp.uri import DataObjectURI, DataspaceURI
 from resqml_objects import parse_resqml_v201_object, serialize_resqml_v201_object
 
 logger = logging.getLogger(__name__)
 
-try:
-    # for py >3.11, we can raise grouped exceptions
-    from builtins import ExceptionGroup  # type: ignore
-except ImportError:
-    # Python 3.10
-    def ExceptionGroup(msg, errors):
-        return errors[0]
-
-
-try:
-    # Python >= 3.11
-    from asyncio import timeout
-except ImportError:
-    # Python 3.10
-    from contextlib import asynccontextmanager
-
-    import async_timeout
-
-    @asynccontextmanager
-    async def timeout(delay: T.Optional[float]) -> T.Any:
-        try:
-            async with async_timeout.timeout(delay):
-                yield None
-        except asyncio.CancelledError as e:
-            raise asyncio.TimeoutError(f"Timeout ({delay}s)") from e
-
-    TimeoutError = asyncio.TimeoutError
-
-try:
-    # Python >= 3.11
-    from typing import Self
-except ImportError:
-    Self = "ETPClient"
-
 
 class ETPError(Exception):
     def __init__(self, message: str, code: int):
@@ -258,11 +174,11 @@ class ETPClient(ETPConnection):
     # client
     #
 
-    async def send(self, body: ETPModel):
+    async def send(self, body: ETPModel) -> list[ETPModel]:
         correlation_id = await self._send(body)
         return await self._recv(correlation_id)
 
-    async def _send(self, body: ETPModel):
+    async def _send(self, body: ETPModel) -> int:
         msg = Message.get_object_message(body, message_flags=MessageFlags.FINALPART)
         if msg is None:
             raise TypeError(f"{type(body)} not valid etp protocol")
@@ -281,41 +197,15 @@ class ETPClient(ETPConnection):
 
         return msg.header.message_id
 
-    async def _recv(self, correlation_id: int) -> ETPModel:
+    async def _recv(self, correlation_id: int) -> list[ETPModel]:
         assert correlation_id in self._recv_events, (
             "Trying to receive a response on non-existing message"
         )
 
-        def timeout_intervals(etp_timeout):
-            # Local function generating progressively longer timeout intervals.
-
-            # Use the timeout-interval generator from the Python websockets
-            # library.
-            backoff_generator = websockets.client.backoff(
-                initial_delay=5.0, min_delay=5.0, max_delay=20.0
-            )
-
-            # Check if we should never time out.
-            if etp_timeout is None:
-                # This is an infinite generator, so it should never exit.
-                yield from backoff_generator
-                return
-
-            # Generate timeout intervals until we have reached the
-            # `etp_timeout`-threshold.
-            csum = 0.0
-            for d in backoff_generator:
-                yield d
-
-                csum += d
-
-                if csum >= etp_timeout:
-                    break
-
         for ti in timeout_intervals(self.etp_timeout):
             try:
                 # Wait for an event for `ti` seconds.
-                async with timeout(ti):
+                async with asyncio.timeout(ti):
                     await self._recv_events[correlation_id].wait()
             except TimeoutError:
                 # Check if the receiver task is still running.
@@ -323,11 +213,18 @@ class ETPClient(ETPConnection):
                     # Raise any errors by waiting for the task to finish.
                     await self.__recvtask
 
-                    logger.error(
-                        "Receiver task terminated without errors. This should not happen"
+                    # Check that the receiver task stopped due to a
+                    # (successfully) closed websockets connection.
+                    try:
+                        await self.ws.recv()
+                    except websockets.ConnectionClosedOK:
+                        pass
+
+                    # Terminate client with an error.
+                    raise ReceiveWorkerExited(
+                        "Receiver task terminated prematurely due to a closed "
+                        "websockets connection"
                     )
-
-                    raise ReceiveWorkerExited
             else:
                 # Break out of for-loop, and start processing message.
                 break
@@ -355,11 +252,7 @@ class ETPClient(ETPConnection):
                 "Server responded with ETPErrors:", ETPError.from_protos(errors)
             )
 
-        if len(bodies) > 1:
-            logger.warning(f"Recived {len(bodies)} messages, but only expected one")
-
-        # ok
-        return bodies[0]
+        return bodies
 
     @staticmethod
     def _parse_error_info(bodies: list[ETPModel]) -> list[ErrorInfo]:
@@ -372,13 +265,15 @@ class ETPClient(ETPConnection):
                 errors.extend(body.errors.values())
         return errors
 
-    async def __aexit__(self, *exc_details) -> None:
-        await self.close(reason="Client exiting")
-
-    async def close(self, reason=""):
+    async def __aexit__(
+        self,
+        exc_type: T.Type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
         close_session_sent = False
         try:
-            await self._send(CloseSession(reason=reason))
+            await self._send(CloseSession(reason="Client exiting"))
             close_session_sent = True
         except websockets.ConnectionClosed:
             logger.error(
@@ -419,7 +314,7 @@ class ETPClient(ETPConnection):
             # In some cases the server does not drop the connection after we
             # have sent the `CloseSession`-message. We therefore add a timeout
             # to the reading of possibly lost messages.
-            async with timeout(self.etp_timeout or 10):
+            async with asyncio.timeout(self.etp_timeout or 10):
                 async for msg in self.ws:
                     counter += 1
         except websockets.ConnectionClosed:
@@ -428,7 +323,7 @@ class ETPClient(ETPConnection):
             pass
         except TimeoutError:
             if close_session_sent:
-                logger.error(
+                logger.warning(
                     "Websockets connection was not closed within "
                     f"{self.etp_timeout or 10} seconds after the "
                     "`CloseSession`-message was sent"
@@ -442,17 +337,30 @@ class ETPClient(ETPConnection):
 
         logger.debug("Client closed")
 
+    async def close(self) -> None:
+        """Closing method that tears down the ETP-connection via the
+        `ETPClient.__aexit__`-method, and closes the websockets connection.
+        This method should _only_ be used if the user has set up a connection
+        via `etp_client = await connect(...)` or `etp_client = await
+        etp_connect(...)` and will handle the closing of the connection
+        manually.
+        """
+
+        await self.__aexit__(None, None, None)
+        # The websockets connection should be closed from the ETP-server once
+        # it has received a `CloseSession`-message. However, calling close on
+        # the websockets connection does not do anything if it is already
+        # closed.
+        await self.ws.close()
+
     async def __recv(self):
         logger.debug("Starting receiver loop")
 
-        while True:
-            # We use this way of receiving messages, instead of the `async
-            # for`-pattern, in order to raise all
-            # `websockets.exceptions.ConnectionClosed`-errors (including the
-            # `websockets.exceptions.ConnectionClosedOK` error). In the `async
-            # for`-case a closing code of `1000` (normal closing) just exits
-            # the loop.
-            msg_data = await self.ws.recv()
+        # Using `async for` makes the receiver task exit without errors on a
+        # `websockets.exceptions.ConnectionClosedOK`-exception. This ensures a
+        # smoother clean-up in case the main-task errors resulting in a closed
+        # websockets connection down the line.
+        async for msg_data in self.ws:
             msg = Message.decode_binary_message(
                 T.cast(bytes, msg_data), ETPClient.generic_transition_table
             )
@@ -468,11 +376,9 @@ class ETPClient(ETPConnection):
                 # set response on send event
                 self._recv_events[msg.header.correlation_id].set()
 
-    #
-    # session related
-    #
+        logger.info("Websockets connection closed and receiver task stopped")
 
-    async def __aenter__(self) -> Self:
+    async def __aenter__(self) -> T.Self:
         return await self.request_session()
 
     async def request_session(self):
@@ -488,7 +394,7 @@ class ETPClient(ETPConnection):
 
             return "store"
 
-        msg = await self.send(
+        msgs = await self.send(
             RequestSession(
                 applicationName=self.application_name,
                 applicationVersion=self.application_version,
@@ -512,6 +418,10 @@ class ETPClient(ETPConnection):
                 ),
             )
         )
+
+        assert len(msgs) == 1
+        msg = msgs[0]
+
         assert msg and isinstance(msg, OpenSession)
 
         self.is_connected = True
@@ -525,12 +435,15 @@ class ETPClient(ETPConnection):
     async def authorize(
         self, authorization: str, supplemental_authorization: T.Mapping[str, str] = {}
     ):
-        msg = await self.send(
+        msgs = await self.send(
             Authorize(
                 authorization=authorization,
                 supplementalAuthorization=supplemental_authorization,
             )
         )
+        assert len(msgs) == 1
+        msg = msgs[0]
+
         assert msg and isinstance(msg, AuthorizeResponse)
 
         return msg
@@ -562,8 +475,10 @@ class ETPClient(ETPConnection):
             raise Exception("Max one / in dataspace name")
         return DataspaceURI.from_name(ds)
 
-    def list_objects(self, dataspace_uri: DataspaceURI | str, depth: int = 1) -> list:
-        return self.send(
+    async def list_objects(
+        self, dataspace_uri: DataspaceURI | str, depth: int = 1
+    ) -> GetResourcesResponse:
+        responses = await self.send(
             GetResources(
                 scope=ContextScopeKind.TARGETS_OR_SELF,
                 context=ContextInfo(
@@ -574,6 +489,8 @@ class ETPClient(ETPConnection):
                 ),
             )
         )
+        assert len(responses) == 1
+        return responses[0]
 
     #
     # dataspace
@@ -582,10 +499,17 @@ class ETPClient(ETPConnection):
     async def get_dataspaces(
         self, store_last_write_filter: int = None
     ) -> GetDataspacesResponse:
-        return await self.send(
+        responses = await self.send(
             GetDataspaces(store_last_write_filter=store_last_write_filter)
         )
 
+        assert all(
+            [isinstance(response, GetDataspacesResponse) for response in responses]
+        ), "Expected GetDataspacesResponse"
+        assert len(responses) == 1
+
+        return responses[0]
+
     async def put_dataspaces(
         self,
         legaltags: list[str],
@@ -593,13 +517,13 @@ class ETPClient(ETPConnection):
         owners: list[str],
         viewers: list[str],
         *dataspace_uris: DataspaceURI,
-    ):
+    ) -> dict[str, str]:
         _uris = list(map(DataspaceURI.from_any, dataspace_uris))
         for i in _uris:
             if i.raw_uri.count("/") > 4:  # includes the 3 eml
                 raise Exception("Max one / in dataspace name")
         time = self.timestamp
-        response = await self.send(
+        responses = await self.send(
             PutDataspaces(
                 dataspaces={
                     d.raw_uri: Dataspace(
@@ -622,15 +546,19 @@ class ETPClient(ETPConnection):
                 }
             )
         )
-        assert isinstance(response, PutDataspacesResponse), (
-            "Expected PutDataspacesResponse"
-        )
+        assert all(
+            [isinstance(response, PutDataspacesResponse) for response in responses]
+        ), "Expected PutDataspacesResponse"
+
+        successes = {}
+        for response in responses:
+            successes = {**successes, **response.success}
 
-        assert len(response.success) == len(dataspace_uris), (
-            f"expected {len(dataspace_uris)} success's"
+        assert len(successes) == len(dataspace_uris), (
+            f"expected {len(dataspace_uris)} successes"
         )
 
-        return response.success
+        return successes
 
     async def put_dataspaces_no_raise(
         self,
@@ -639,22 +567,30 @@ class ETPClient(ETPConnection):
         owners: list[str],
         viewers: list[str],
         *dataspace_uris: DataspaceURI,
-    ):
+    ) -> dict[str, str]:
         try:
             return await self.put_dataspaces(
                 legaltags, otherRelevantDataCountries, owners, viewers, *dataspace_uris
             )
         except ETPError:
-            pass
+            return {}
 
-    async def delete_dataspaces(self, *dataspace_uris: DataspaceURI):
+    async def delete_dataspaces(self, *dataspace_uris: DataspaceURI) -> dict[str, str]:
         _uris = list(map(str, dataspace_uris))
 
-        response = await self.send(DeleteDataspaces(uris=dict(zip(_uris, _uris))))
-        assert isinstance(response, DeleteDataspacesResponse), (
-            "Expected DeleteDataspacesResponse"
+        responses = await self.send(DeleteDataspaces(uris=dict(zip(_uris, _uris))))
+        assert all(
+            [isinstance(response, DeleteDataspacesResponse) for response in responses]
+        ), "Expected DeleteDataspacesResponse"
+
+        successes = {}
+        for response in responses:
+            successes = {**successes, **response.success}
+
+        assert len(successes) == len(dataspace_uris), (
+            f"expected {len(dataspace_uris)} successes"
         )
-        return response.success
+        return successes
 
     async def get_data_objects(self, *uris: T.Union[DataObjectURI, str]):
         tasks = []
@@ -662,7 +598,8 @@ class ETPClient(ETPConnection):
             task = self.send(GetDataObjects(uris={str(uri): str(uri)}))
             tasks.append(task)
 
-        responses = await asyncio.gather(*tasks)
+        task_responses = await asyncio.gather(*tasks)
+        responses = [r for tr in task_responses for r in tr]
         assert len(responses) == len(uris)
 
         data_objects = []
@@ -695,7 +632,8 @@ class ETPClient(ETPConnection):
             )
             tasks.append(task)
 
-        responses = await asyncio.gather(*tasks)
+        task_responses = await asyncio.gather(*tasks)
+        responses = [r for tr in task_responses for r in tr]
 
         errors = []
         for response in responses:
@@ -757,12 +695,16 @@ class ETPClient(ETPConnection):
     ):
         _uris = list(map(str, uris))
 
-        response = await self.send(
+        responses = await self.send(
             DeleteDataObjects(
                 uris=dict(zip(_uris, _uris)),
                 prune_contained_objects=prune_contained_objects,
             )
         )
+
+        assert len(responses) == 1
+        response = responses[0]
+
         assert isinstance(response, DeleteDataObjectsResponse), (
             "Expected DeleteDataObjectsResponse"
         )
@@ -770,23 +712,34 @@ class ETPClient(ETPConnection):
         return response.deleted_uris
 
     async def start_transaction(
-        self, dataspace_uri: DataspaceURI, read_only: bool = True
+        self, dataspace_uri: DataspaceURI | str, read_only: bool = True
     ) -> Uuid:
-        trans_id = await self.send(
-            StartTransaction(
-                read_only=read_only, dataspace_uris=[dataspace_uri.raw_uri]
-            )
+        dataspace_uri = str(DataspaceURI.from_any(dataspace_uri))
+        responses = await self.send(
+            StartTransaction(read_only=read_only, dataspace_uris=[dataspace_uri])
         )
-        if trans_id.successful is False:
-            raise Exception(f"Failed starting transaction {dataspace_uri.raw_uri}")
-        # uuid.UUID(bytes=trans_id.transaction_uuid)
-        return Uuid(trans_id.transaction_uuid)
+        assert all(
+            [isinstance(response, StartTransactionResponse) for response in responses]
+        ), "Expected StartTransactionResponse"
+
+        assert len(responses) == 1
+        response = responses[0]
+
+        if not response.successful:
+            raise ETPTransactionFailure(f"Failed starting transaction {dataspace_uri}")
+
+        return response.transaction_uuid
 
     async def commit_transaction(self, transaction_uuid: Uuid):
-        r = await self.send(CommitTransaction(transaction_uuid=transaction_uuid))
-        if r.successful is False:
-            raise Exception(r.failure_reason)
-        return r
+        responses = await self.send(
+            CommitTransaction(transaction_uuid=transaction_uuid)
+        )
+        assert len(responses) == 1
+        response = responses[0]
+
+        if response.successful is False:
+            raise ETPTransactionFailure(response.failure_reason)
+        return response
 
     async def rollback_transaction(self, transaction_id: Uuid):
         return await self.send(RollbackTransaction(transactionUuid=transaction_id))
@@ -811,20 +764,19 @@ class ETPClient(ETPConnection):
             "obj must be Grid2DRepresentation"
         )
         sgeo = gri.grid2d_patch.geometry.points.supporting_geometry  # type: ignore
-        if sys.version_info[1] != 10:
-            assert isinstance(
-                gri.grid2d_patch.geometry.points, ro.Point3dZValueArray
-            ), "Points must be Point3dZValueArray"
-            assert isinstance(
-                gri.grid2d_patch.geometry.points.zvalues, ro.DoubleHdf5Array
-            ), "Values must be DoubleHdf5Array"
-            assert isinstance(
-                gri.grid2d_patch.geometry.points.supporting_geometry,
-                ro.Point3dLatticeArray,
-            ), "Points support_geo must be Point3dLatticeArray"
-            assert isinstance(sgeo, ro.Point3dLatticeArray), (
-                "supporting_geometry must be Point3dLatticeArray"
-            )
+        assert isinstance(gri.grid2d_patch.geometry.points, ro.Point3dZValueArray), (
+            "Points must be Point3dZValueArray"
+        )
+        assert isinstance(
+            gri.grid2d_patch.geometry.points.zvalues, ro.DoubleHdf5Array
+        ), "Values must be DoubleHdf5Array"
+        assert isinstance(
+            gri.grid2d_patch.geometry.points.supporting_geometry,
+            ro.Point3dLatticeArray,
+        ), "Points support_geo must be Point3dLatticeArray"
+        assert isinstance(sgeo, ro.Point3dLatticeArray), (
+            "supporting_geometry must be Point3dLatticeArray"
+        )
         assert isinstance(
             gri.grid2d_patch.geometry.points.zvalues.values, ro.Hdf5Dataset
         ), "Values must be Hdf5Dataset"
@@ -893,9 +845,12 @@ class ETPClient(ETPConnection):
     #
 
     async def get_array_metadata(self, *uids: DataArrayIdentifier):
-        response = await self.send(
+        responses = await self.send(
             GetDataArrayMetadata(dataArrays={i.path_in_resource: i for i in uids})
         )
+        assert len(responses) == 1
+        response = responses[0]
+
         assert isinstance(response, GetDataArrayMetadataResponse)
 
         if len(response.array_metadata) != len(uids):
@@ -915,9 +870,13 @@ class ETPClient(ETPConnection):
         ):
             return await self._get_array_chunked(uid)
 
-        response = await self.send(
+        responses = await self.send(
             GetDataArrays(dataArrays={uid.path_in_resource: uid})
         )
+
+        assert len(responses) == 1
+        response = responses[0]
+
         assert isinstance(response, GetDataArraysResponse), (
             "Expected GetDataArraysResponse"
         )
@@ -936,10 +895,13 @@ class ETPClient(ETPConnection):
             path_in_resource=path_in_resource,
         )
 
-        response = await self.send(
+        responses = await self.send(
             GetDataArrayMetadata(data_arrays={dai.path_in_resource: dai}),
         )
 
+        assert len(responses) == 1
+        response = responses[0]
+
         self.assert_response(response, GetDataArrayMetadataResponse)
         assert (
             len(response.array_metadata) == 1
@@ -1002,7 +964,12 @@ class ETPClient(ETPConnection):
                 )
                 tasks.append(task)
 
-            responses = await asyncio.gather(*tasks)
+            task_responses = await asyncio.gather(*tasks)
+            responses = [
+                response
+                for task_response in task_responses
+                for response in task_response
+            ]
 
             data_blocks = []
             for i, response in enumerate(responses):
@@ -1037,10 +1004,13 @@ class ETPClient(ETPConnection):
             return data
 
         # Download the full array in one go.
-        response = await self.send(
+        responses = await self.send(
             GetDataArrays(data_arrays={dai.path_in_resource: dai}),
         )
 
+        assert len(responses) == 1
+        response = responses[0]
+
         self.assert_response(response, GetDataArraysResponse)
         assert (
             len(response.data_arrays) == 1
@@ -1072,7 +1042,7 @@ class ETPClient(ETPConnection):
         now = self.timestamp
 
         # Allocate space on server for the array.
-        response = await self.send(
+        responses = await self.send(
             PutUninitializedDataArrays(
                 data_arrays={
                     dai.path_in_resource: PutUninitializedDataArrayType(
@@ -1089,6 +1059,9 @@ class ETPClient(ETPConnection):
             ),
         )
 
+        assert len(responses) == 1
+        response = responses[0]
+
         self.assert_response(response, PutUninitializedDataArraysResponse)
         assert len(response.success) == 1 and dai.path_in_resource in response.success
 
@@ -1138,7 +1111,14 @@ class ETPClient(ETPConnection):
                 tasks.append(task)
 
             # Upload all blocks.
-            responses = await asyncio.gather(*tasks)
+            task_responses = await asyncio.gather(*tasks)
+
+            # Flatten list of responses.
+            responses = [
+                response
+                for task_response in task_responses
+                for response in task_response
+            ]
 
             # Check for successful responses.
             for response in responses:
@@ -1155,7 +1135,7 @@ class ETPClient(ETPConnection):
         etp_array_data = utils_arrays.get_etp_data_array_from_numpy(data)
 
         # Pass entire array in one message.
-        response = await self.send(
+        responses = await self.send(
             PutDataArrays(
                 data_arrays={
                     dai.path_in_resource: PutDataArraysType(
@@ -1166,6 +1146,9 @@ class ETPClient(ETPConnection):
             )
         )
 
+        assert len(responses) == 1
+        response = responses[0]
+
         self.assert_response(response, PutDataArraysResponse)
         assert len(response.success) == 1 and dai.path_in_resource in response.success
 
@@ -1187,7 +1170,7 @@ class ETPClient(ETPConnection):
         if data.nbytes > self.max_array_size:
             return await self._put_array_chunked(uid, data)
 
-        response = await self.send(
+        responses = await self.send(
             PutDataArrays(
                 data_arrays={
                     uid.path_in_resource: PutDataArraysType(
@@ -1198,6 +1181,9 @@ class ETPClient(ETPConnection):
             )
         )
 
+        assert len(responses) == 1
+        response = responses[0]
+
         assert isinstance(response, PutDataArraysResponse), (
             "Expected PutDataArraysResponse"
         )
@@ -1221,9 +1207,13 @@ class ETPClient(ETPConnection):
             starts=starts.tolist(),
             counts=counts.tolist(),
         )
-        response = await self.send(
+        responses = await self.send(
             GetDataSubarrays(dataSubarrays={uid.path_in_resource: payload})
         )
+
+        assert len(responses) == 1
+        response = responses[0]
+
         assert isinstance(response, GetDataSubarraysResponse), (
             "Expected GetDataSubarraysResponse"
         )
@@ -1262,9 +1252,13 @@ class ETPClient(ETPConnection):
             f"{dataarray.data.item.__class__.__name__}"
         )
 
-        response = await self.send(
+        responses = await self.send(
             PutDataSubarrays(dataSubarrays={uid.path_in_resource: payload})
         )
+
+        assert len(responses) == 1
+        response = responses[0]
+
         assert isinstance(response, PutDataSubarraysResponse), (
             "Expected PutDataSubarraysResponse"
         )
@@ -1386,9 +1380,13 @@ class ETPClient(ETPConnection):
                 )
             ),
         )
-        response = await self.send(
+        responses = await self.send(
             PutUninitializedDataArrays(dataArrays={uid.path_in_resource: payload})
         )
+
+        assert len(responses) == 1
+        response = responses[0]
+
         assert isinstance(response, PutUninitializedDataArraysResponse), (
             "Expected PutUninitializedDataArraysResponse"
         )
@@ -1407,6 +1405,7 @@ class connect:
     # ... = await connect(...)
 
     def __await__(self):
+        # The caller is response for calling `close()` on the client.
         return self.__aenter__().__await__()
 
     # async with connect(...) as ...:
@@ -1420,7 +1419,7 @@ class connect:
         if self.data_partition is not None:
             headers["data-partition-id"] = self.data_partition
 
-        self.ws = await websockets.connect(
+        ws = await websockets.connect(
             self.server_url,
             subprotocols=[ETPClient.SUB_PROTOCOL],  # type: ignore
             additional_headers=headers,
@@ -1430,7 +1429,7 @@ class connect:
         )
 
         self.client = ETPClient(
-            self.ws,
+            ws=ws,
             etp_timeout=self.timeout,
             max_message_size=SETTINGS.MaxWebSocketMessagePayloadSize,
             application_name=SETTINGS.application_name,
@@ -1448,67 +1447,188 @@ class connect:
 
     # exit the async context manager
     async def __aexit__(self, exc_type, exc: Exception, tb: TracebackType):
+        # The `ETPClient.close`-method also closes the websockets connection.
         await self.client.close()
-        await self.ws.close()
 
 
-@contextlib.asynccontextmanager
-async def etp_connect(
-    uri: str,
-    data_partition_id: str | None = None,
-    authorization: str | None = None,
-    etp_timeout: float = 10.0,
-    max_message_size: float = 2**20,
-) -> ETPClient:
-    additional_headers = {}
-
-    if authorization is not None:
-        additional_headers["Authorization"] = authorization
-    if data_partition_id is not None:
-        additional_headers["data-partition-id"] = data_partition_id
-
-    subprotocols = ["etp12.energistics.org"]
-
-    async with (
-        websockets.connect(
-            uri=uri,
-            subprotocols=subprotocols,
-            max_size=max_message_size,
-            additional_headers=additional_headers,
-        ) as ws,
-        ETPClient(
-            ws=ws,
-            etp_timeout=etp_timeout,
-            max_message_size=max_message_size,
-        ) as etp_client,
-    ):
-        yield etp_client
-
-
-async def etp_persistent_connect(
-    uri: str,
-    data_partition_id: str | None = None,
-    authorization: str | None = None,
-    etp_timeout: float = 10.0,
-    max_message_size: float = 2**20,
-) -> AsyncGenerator[ETPClient]:
-    additional_headers = {}
-
-    if authorization is not None:
-        additional_headers["Authorization"] = authorization
-    if data_partition_id is not None:
-        additional_headers["data-partition-id"] = data_partition_id
-
-    subprotocols = ["etp12.energistics.org"]
-    async for ws in websockets.connect(
-        uri=uri,
-        subprotocols=subprotocols,
-        max_size=max_message_size,
-        additional_headers=additional_headers,
-    ):
-        async with ETPClient(
-            ws=ws,
-            etp_timeout=etp_timeout,
-            max_message_size=max_message_size,
-        ) as etp_client:
-            yield etp_client
+class etp_connect:
+    """
+    Connect to an ETP server via websockets.
+
+    This class can act as:
+
+    1. A context manager handling setup and tear-down of the connection.
+    2. An asynchronous iterator which can be used to persistently retry to
+    connect if the websockets connection drops.
+    3. An awaitable connection that must be manually closed by the user.
+
+    See below for examples of all three cases.
+
+    Parameters
+    ----------
+    uri: str
+        The uri to the ETP server. This should be the uri to a websockets
+        endpoint.
+    data_partition_id: str | None
+        The data partition id used when connecting to the OSDU open-etp-server
+        in multi-partition mode. Default is `None`.
+    authorization: str | SecretStr | None
+        Bearer token used for authenticating to the ETP server. This token
+        should be on the form `"Bearer 1234..."`. Default is `None`.
+    etp_timeout: float | None
+        The timeout in seconds for when to stop waiting for a message from the
+        ETP server. Setting it to `None` will persist the connection
+        indefinetly. Default is `None`.
+    max_message_size: float
+        The maximum number of bytes for a single websockets message. Default is
+        `2**20` corresponding to `1` MiB.
+
+
+    Examples
+    --------
+    An example of connecting to the ETP server using :func:`etp_connect` as a
+    context manager is:
+
+        async with etp_connect(...) as etp_client:
+            ...
+
+    In this case the closing message and the websockets connection is closed
+    once the program exits the context manager.
+
+
+    To persist a connection if the websockets connection is dropped (for any
+    reason), use :func:`etp_connect` as an asynchronous generator, viz.:
+
+        import websockets
+
+        async for etp_client in etp_connect(...):
+            try:
+                ...
+            except websockets.ConnectionClosed:
+                continue
+
+            # Include `break` to avoid re-running the whole block if the
+            # iteration runs without any errors.
+            break
+
+    Note that in this case the whole program under the `try`-block is re-run
+    from the start if the iteration completes normally, or if the websockets
+    connection is dropped. Therefore, make sure to include a `break` at the end
+    of the `try`-block (as in the example above).
+
+
+    The third option is to set up a connection via `await` and then manually
+    close the connection once done:
+
+        etp_client = await etp_connect(...)
+        ...
+        await etp_client.close()
+    """
+
+    def __init__(
+        self,
+        uri: str,
+        data_partition_id: str | None = None,
+        authorization: str | SecretStr | None = None,
+        etp_timeout: float | None = None,
+        max_message_size: float = 2**20,
+    ) -> None:
+        self.uri = uri
+        self.data_partition_id = data_partition_id
+
+        if isinstance(authorization, SecretStr):
+            self.authorization = authorization
+        else:
+            self.authorization = SecretStr(authorization)
+
+        self.etp_timeout = etp_timeout
+        self.max_message_size = max_message_size
+        self.subprotocols = ["etp12.energistics.org"]
+
+    def __await__(self) -> ETPClient:
+        # The caller is responsible for calling `close()` on the client.
+        return self.__aenter__().__await__()
+
+    def get_additional_headers(self) -> dict[str, str]:
+        additional_headers = {}
+
+        if self.authorization.get_secret_value() is not None:
+            additional_headers["Authorization"] = self.authorization.get_secret_value()
+
+        if self.data_partition_id is not None:
+            additional_headers["data-partition-id"] = self.data_partition_id
+
+        return additional_headers
+
+    async def __aenter__(self) -> ETPClient:
+        self.stack = contextlib.AsyncExitStack()
+        try:
+            ws = await self.stack.enter_async_context(
+                websockets.connect(
+                    uri=self.uri,
+                    subprotocols=self.subprotocols,
+                    max_size=self.max_message_size,
+                    additional_headers=self.get_additional_headers(),
+                )
+            )
+            etp_client = await self.stack.enter_async_context(
+                ETPClient(
+                    ws=ws,
+                    etp_timeout=self.etp_timeout,
+                    max_message_size=self.max_message_size,
+                )
+            )
+        except BaseException:
+            await self.stack.aclose()
+            raise
+
+        return etp_client
+
+    async def __aexit__(
+        self,
+        exc_type: T.Type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        return await self.stack.aclose()
+
+    async def __aiter__(self) -> AsyncGenerator[ETPClient]:
+        async for ws in websockets.connect(
+            uri=self.uri,
+            subprotocols=self.subprotocols,
+            max_size=self.max_message_size,
+            additional_headers=self.get_additional_headers(),
+        ):
+            async with ETPClient(
+                ws=ws,
+                etp_timeout=self.etp_timeout,
+                max_message_size=self.max_message_size,
+            ) as etp_client:
+                yield etp_client
+
+
+def timeout_intervals(total_timeout: float) -> Generator[float]:
+    # Local function generating progressively longer timeout intervals.
+
+    # Use the timeout-interval generator from the Python websockets
+    # library.
+    backoff_generator = websockets.client.backoff(
+        initial_delay=5.0, min_delay=5.0, max_delay=20.0
+    )
+
+    # Check if we should never time out.
+    if total_timeout is None:
+        # This is an infinite generator, so it should never exit.
+        yield from backoff_generator
+        return
+
+    # Generate timeout intervals until we have reached the
+    # `total_timeout`-threshold.
+    csum = 0.0
+    for d in backoff_generator:
+        yield d
+
+        csum += d
+
+        if csum >= total_timeout:
+            break