From f73896a244146f43d9188d286c2871052b53c872 Mon Sep 17 00:00:00 2001
From: Jonas Maison <Jonas1312@users.noreply.github.com>
Date: Thu, 5 Oct 2023 08:52:16 +0200
Subject: [PATCH] refac:

---
 src/kili/domain/cloud_storage.py              |  3 +-
 src/kili/presentation/client/cloud_storage.py | 39 +++++++++++--------
 src/kili/use_cases/cloud_storage/__init__.py  | 17 ++------
 3 files changed, 28 insertions(+), 31 deletions(-)

diff --git a/src/kili/domain/cloud_storage.py b/src/kili/domain/cloud_storage.py
index d4d4576a2..fde9e400a 100644
--- a/src/kili/domain/cloud_storage.py
+++ b/src/kili/domain/cloud_storage.py
@@ -35,8 +35,7 @@ class DataIntegrationFilters:
 
 @dataclass
 class DataConnectionFilters:
-    """Data Connection filters."""
+    """DataConnectionsWhere filters."""
 
-    data_connection_id: Optional[DataConnectionId]
     project_id: Optional[ProjectId]
     integration_id: Optional[DataIntegrationId]
diff --git a/src/kili/presentation/client/cloud_storage.py b/src/kili/presentation/client/cloud_storage.py
index 7871c02cc..ab7195ca1 100644
--- a/src/kili/presentation/client/cloud_storage.py
+++ b/src/kili/presentation/client/cloud_storage.py
@@ -127,23 +127,30 @@ def cloud_storage_connections(
 
         disable_tqdm = disable_tqdm_if_as_generator(as_generator, disable_tqdm)
 
-        data_connections_gen = CloudStorageUseCases(self.kili_api_gateway).list_data_connections(
-            data_connection_filters=DataConnectionFilters(
-                data_connection_id=(
-                    DataConnectionId(cloud_storage_connection_id)
-                    if cloud_storage_connection_id is not None
-                    else None
-                ),
-                project_id=ProjectId(project_id) if project_id is not None else None,
-                integration_id=(
-                    DataIntegrationId(cloud_storage_integration_id)
-                    if cloud_storage_integration_id is not None
-                    else None
+        cloud_storage_use_cases = CloudStorageUseCases(self.kili_api_gateway)
+
+        if cloud_storage_connection_id is None:
+            data_connections_gen = cloud_storage_use_cases.list_data_connections(
+                data_connection_filters=DataConnectionFilters(
+                    project_id=ProjectId(project_id) if project_id is not None else None,
+                    integration_id=(
+                        DataIntegrationId(cloud_storage_integration_id)
+                        if cloud_storage_integration_id is not None
+                        else None
+                    ),
                 ),
-            ),
-            fields=fields,
-            options=QueryOptions(disable_tqdm, first, skip),
-        )
+                fields=fields,
+                options=QueryOptions(disable_tqdm, first, skip),
+            )
+        else:
+            data_connections_gen = (
+                i
+                for i in [
+                    cloud_storage_use_cases.get_data_connection(
+                        DataConnectionId(cloud_storage_connection_id), fields=fields
+                    )
+                ]
+            )
 
         if as_generator:
             return data_connections_gen
diff --git a/src/kili/use_cases/cloud_storage/__init__.py b/src/kili/use_cases/cloud_storage/__init__.py
index 8c0673256..e575645ed 100644
--- a/src/kili/use_cases/cloud_storage/__init__.py
+++ b/src/kili/use_cases/cloud_storage/__init__.py
@@ -1,7 +1,7 @@
 """Cloud storage use cases."""
 import time
 from datetime import datetime
-from typing import Dict, Generator, Iterable, List, Optional
+from typing import Dict, Generator, List, Optional
 
 from tenacity import Retrying
 from tenacity.retry import retry_if_exception_type
@@ -34,19 +34,10 @@ def list_data_connections(
         data_connection_filters: DataConnectionFilters,
         fields: ListOrTuple[str],
         options: QueryOptions,
-    ) -> Iterable[Dict]:
+    ) -> Generator[Dict, None, None]:
         """List data connections."""
-        if data_connection_filters.data_connection_id is None:
-            return self._kili_api_gateway.list_data_connections(
-                data_connection_filters=data_connection_filters, fields=fields, options=options
-            )
-
-        return iter(
-            [
-                self._kili_api_gateway.get_data_connection(
-                    data_connection_id=data_connection_filters.data_connection_id, fields=fields
-                )
-            ]
+        return self._kili_api_gateway.list_data_connections(
+            data_connection_filters=data_connection_filters, fields=fields, options=options
         )
 
     def get_data_connection(