fixing test

iterative · Aug 15, 2024 · 1df4c26 · 1df4c26
1 parent 73f291a
commit 1df4c26
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/src/datachain/lib/dc.py b/src/datachain/lib/dc.py
@@ -386,7 +386,10 @@ def from_storage(
             client_config["anon"] = True
 
         session = Session.get(
-            session, catalog=kwargs.get("catalog"), client_config=client_config
+            session,
+            catalog=kwargs.get("catalog"),
+            client_config=client_config,
+            in_memory=in_memory,
         )
 
         client, path = Client.parse_url(
@@ -404,7 +407,7 @@ def from_storage(
             client.uri, posixpath.join(lst_path, "")
         )
 
-        for ds in cls.datasets(session=session).collect("dataset"):
+        for ds in cls.datasets(session=session, in_memory=in_memory).collect("dataset"):
             # TODO filter out expired ones
             if (
                 DatasetRecord.contains_listing(
@@ -421,7 +424,12 @@ def from_storage(
                 )
 
         (
-            cls.from_records(DataChain.DEFAULT_FILE_RECORD, session=session, **kwargs)
+            cls.from_records(
+                DataChain.DEFAULT_FILE_RECORD,
+                session=session,
+                in_memory=in_memory,
+                **kwargs,
+            )
             .gen(
                 list_bucket(lst_uri, **session.catalog.client_config),
                 output={f"{object_name}": file_type},