Lal/dataset api v5 prod 14582

cosmotech/coal/cosmotech_api/__init__.py

@@ -17,7 +17,8 @@
 )
 
 from cosmotech.coal.utils.semver import semver_of
-csm_version = semver_of('cosmotech_api')
+
+csm_version = semver_of("cosmotech_api")
 if csm_version.major < 5:
     # Re-export functions from the twin_data_layer module
     from cosmotech.coal.cosmotech_api.twin_data_layer import (

cosmotech/coal/cosmotech_api/dataset/download/file.py

@@ -52,9 +52,7 @@ def process_xls(target_file) -> Dict[str, Any]:
                 content[sheet_name].append(new_row)
                 row_count += 1
 
-        LOGGER.debug(
-            T("coal.services.dataset.sheet_processed").format(sheet_name=sheet_name, rows=row_count)
-        )
+        LOGGER.debug(T("coal.services.dataset.sheet_processed").format(sheet_name=sheet_name, rows=row_count))
     return content
 
 
@@ -87,9 +85,7 @@ def process_csv(target_file) -> Dict[str, Any]:
             content[current_filename].append(new_row)
             row_count += 1
 
-        LOGGER.debug(
-            T("coal.services.dataset.csv_processed").format(file_name=current_filename, rows=row_count)
-        )
+        LOGGER.debug(T("coal.services.dataset.csv_processed").format(file_name=current_filename, rows=row_count))
     return content
 
 
@@ -107,9 +103,7 @@ def process_json(target_file) -> Dict[str, Any]:
         else:
             item_count = 1
 
-        LOGGER.debug(
-            T("coal.services.dataset.json_processed").format(file_name=current_filename, items=item_count)
-        )
+        LOGGER.debug(T("coal.services.dataset.json_processed").format(file_name=current_filename, items=item_count))
     return content
 
 
@@ -121,9 +115,7 @@ def process_txt(target_file) -> Dict[str, Any]:
         content[current_filename] = "".join(line for line in _file)
 
         line_count = content[current_filename].count("\n") + 1
-        LOGGER.debug(
-            T("coal.services.dataset.text_processed").format(file_name=current_filename, lines=line_count)
-        )
+        LOGGER.debug(T("coal.services.dataset.text_processed").format(file_name=current_filename, lines=line_count))
     return content
 
 
@@ -140,6 +132,7 @@ def timed_read_file(file_name, file):
         else:
             content.update(process_txt(file))
         return content
+
     return timed_read_file(file_name, file)
 
 

cosmotech/coal/cosmotech_api/dataset/download/twingraph.py

@@ -62,7 +62,9 @@ def download_twingraph_dataset(
         # Query edges
         edges_start = time.time()
         LOGGER.debug(T("coal.services.dataset.twingraph_querying_edges").format(dataset_id=dataset_id))
-        edges_query = cosmotech_api.DatasetTwinGraphQuery(query="MATCH(n)-[r]->(m) RETURN n as src, r as rel, m as dest")
+        edges_query = cosmotech_api.DatasetTwinGraphQuery(
+            query="MATCH(n)-[r]->(m) RETURN n as src, r as rel, m as dest"
+        )
 
         edges = dataset_api.twingraph_query(
             organization_id=organization_id,

cosmotech/coal/cosmotech_api/runner/datasets.py

@@ -57,7 +57,7 @@ def download_dataset(
     read_files: bool = True,
 ) -> Dict[str, Any]:
     """
-        retro-compatibility to cosmo-api v4
+    retro-compatibility to cosmo-api v4
     """
     from cosmotech.coal.utils.semver import semver_of
 
@@ -90,9 +90,9 @@ def download_dataset_v5(
     # Get dataset information
     with get_api_client()[0] as api_client:
         dataset_api_instance = DatasetApi(api_client)
-        dataset = dataset_api_instance.get_dataset(organization_id=organization_id,
-                                                   workspace_id=workspace_id,
-                                                   dataset_id=dataset_id)
+        dataset = dataset_api_instance.get_dataset(
+            organization_id=organization_id, workspace_id=workspace_id, dataset_id=dataset_id
+        )
 
         content = dict()
         tmp_dataset_dir = tempfile.mkdtemp()
@@ -223,9 +223,7 @@ def download_dataset_v4(
             }
 
 
-def download_dataset_process(
-    _dataset_id, organization_id, workspace_id, read_files, _return_dict, _error_dict
-):
+def download_dataset_process(_dataset_id, organization_id, workspace_id, read_files, _return_dict, _error_dict):
     """
     Process function for downloading a dataset in a separate process.
 

cosmotech/coal/utils/decorator.py

@@ -19,5 +19,7 @@ def wrapper(*args, **kwargs):
             else:
                 LOGGER.info(msg)
             return r
+
         return wrapper
+
     return decorator

tests/unit/coal/test_cosmotech_api/test_cosmotech_api_run_template.py

@@ -7,8 +7,7 @@
 
 import pathlib
 import pytest
-from unittest.mock import MagicMock, patch, mock_open
-from io import BytesIO
+from unittest.mock import MagicMock, patch
 from zipfile import BadZipfile, ZipFile
 
 import cosmotech_api
@@ -17,8 +16,10 @@
 from cosmotech_api.exceptions import ServiceException
 
 from cosmotech.coal.cosmotech_api.run_template import load_run_template_handlers
+from cosmotech.coal.utils.semver import semver_of
 
 
+@pytest.mark.skipif(semver_of("cosmotech_api").major >= 5, reason="not supported in version 5")
 class TestRunTemplateFunctions:
     """Tests for top-level functions in the run_template module."""
 
@@ -66,7 +67,6 @@ def test_load_run_template_handlers_success(self):
             patch("cosmotech.coal.cosmotech_api.run_template.WorkspaceApi", return_value=mock_workspace_api),
             patch("cosmotech.coal.cosmotech_api.run_template.SolutionApi", return_value=mock_solution_api),
             patch("cosmotech.coal.cosmotech_api.run_template.ZipFile", return_value=mock_zipfile_context),
-            patch("cosmotech.coal.cosmotech_api.run_template.BytesIO") as mock_bytesio,
             patch("cosmotech.coal.cosmotech_api.run_template.pathlib.Path") as mock_path_class,
         ):
             mock_path_class.return_value = mock_path

tests/unit/coal/test_cosmotech_api/test_cosmotech_api_twin_data_layer.py

@@ -14,22 +14,29 @@
 
 import pytest
 import requests
-from cosmotech_api import DatasetApi, RunnerApi, DatasetTwinGraphQuery
-
-from cosmotech.coal.cosmotech_api.twin_data_layer import (
-    get_dataset_id_from_runner,
-    send_files_to_tdl,
-    load_files_from_tdl,
-    CSVSourceFile,
-    _process_csv_file,
-    _get_node_properties,
-    _get_relationship_properties,
-    _execute_queries,
-    _write_files,
-    ID_COLUMN,
-    SOURCE_COLUMN,
-    TARGET_COLUMN,
-)
+
+from cosmotech.coal.utils.semver import semver_of
+from cosmotech_api import DatasetApi, RunnerApi
+
+if semver_of("cosmotech_api").major < 5:
+    from cosmotech_api import DatasetTwinGraphQuery
+
+    from cosmotech.coal.cosmotech_api.twin_data_layer import (
+        get_dataset_id_from_runner,
+        send_files_to_tdl,
+        load_files_from_tdl,
+        CSVSourceFile,
+        _process_csv_file,
+        _get_node_properties,
+        _get_relationship_properties,
+        _execute_queries,
+        _write_files,
+        ID_COLUMN,
+        SOURCE_COLUMN,
+        TARGET_COLUMN,
+    )
+
+pytestmark = pytest.mark.skipif(semver_of("cosmotech_api").major >= 5, reason="not supported under version 5")
 
 
 class TestCSVSourceFile:

tests/unit/coal/test_cosmotech_api/test_cosmotech_api_twin_data_layer_auth.py

@@ -14,18 +14,26 @@
 
 import pytest
 import requests
-from cosmotech_api import DatasetApi, RunnerApi, DatasetTwinGraphQuery
+
+from cosmotech.coal.utils.semver import semver_of
+from cosmotech_api import DatasetApi, RunnerApi
+
+if semver_of("cosmotech_api").major < 5:
+    from cosmotech_api import DatasetTwinGraphQuery
+    from cosmotech.coal.cosmotech_api.twin_data_layer import (
+        send_files_to_tdl,
+        load_files_from_tdl,
+        _process_csv_file,
+        _get_node_properties,
+        _get_relationship_properties,
+    )
 
 from cosmotech.orchestrator.utils.translate import T
-from cosmotech.coal.cosmotech_api.twin_data_layer import (
-    send_files_to_tdl,
-    load_files_from_tdl,
-    _process_csv_file,
-    _get_node_properties,
-    _get_relationship_properties,
-)
+
+skip_under_v5 = pytest.mark.skipif(semver_of("cosmotech_api").major >= 5, reason="not supported under version 5")
 
 
+@skip_under_v5
 class TestTwinDataLayerAuth:
     """Tests for authentication in the twin_data_layer module."""
 

tests/unit/coal/test_cosmotech_api/test_cosmotech_api_twin_data_layer_coverage.py

@@ -17,18 +17,26 @@
 
 import pytest
 import requests
-from cosmotech_api import DatasetApi, RunnerApi, DatasetTwinGraphQuery
 
-from cosmotech.coal.cosmotech_api.twin_data_layer import (
-    get_dataset_id_from_runner,
-    send_files_to_tdl,
-    load_files_from_tdl,
-    _process_csv_file,
-    _write_files,
-    BATCH_SIZE_LIMIT,
-)
+from cosmotech.coal.utils.semver import semver_of
+from cosmotech_api import DatasetApi, RunnerApi
 
+if semver_of("cosmotech_api").major < 5:
+    from cosmotech_api import DatasetTwinGraphQuery
 
+    from cosmotech.coal.cosmotech_api.twin_data_layer import (
+        get_dataset_id_from_runner,
+        send_files_to_tdl,
+        load_files_from_tdl,
+        _process_csv_file,
+        _write_files,
+        BATCH_SIZE_LIMIT,
+    )
+
+skip_under_v5 = pytest.mark.skipif(semver_of("cosmotech_api").major >= 5, reason="not supported under version 5")
+
+
+@skip_under_v5
 class TestTwinDataLayerCoverage:
     """Additional tests for the twin_data_layer module to improve coverage."""
 

tests/unit/coal/test_cosmotech_api/test_cosmotech_api_twin_data_layer_edge_cases.py

@@ -14,18 +14,25 @@
 
 import pytest
 import requests
-from cosmotech_api import DatasetApi, RunnerApi, DatasetTwinGraphQuery
 
+from cosmotech.coal.utils.semver import semver_of
+from cosmotech_api import DatasetApi, RunnerApi
+
+if semver_of("cosmotech_api").major < 5:
+    from cosmotech_api import DatasetTwinGraphQuery
+    from cosmotech.coal.cosmotech_api.twin_data_layer import (
+        send_files_to_tdl,
+        load_files_from_tdl,
+        _process_csv_file,
+        _get_node_properties,
+        _get_relationship_properties,
+    )
 from cosmotech.orchestrator.utils.translate import T
-from cosmotech.coal.cosmotech_api.twin_data_layer import (
-    send_files_to_tdl,
-    load_files_from_tdl,
-    _process_csv_file,
-    _get_node_properties,
-    _get_relationship_properties,
-)
 
+skip_under_v5 = pytest.mark.skipif(semver_of("cosmotech_api").major >= 5, reason="not supported under version 5")
 
+
+@skip_under_v5
 class TestTwinDataLayerEdgeCases:
     """Tests for edge cases in the twin_data_layer module."""
 

tests/unit/coal/test_cosmotech_api/test_dataset/test_download/test_download_common.py

@@ -6,14 +6,16 @@
 # specifically authorized by written means by Cosmo Tech.
 
 import pytest
-from unittest.mock import MagicMock, patch, call
+from unittest.mock import MagicMock, patch
 from pathlib import Path
 
 from cosmotech_api import DatasetApi
 
 from cosmotech.coal.cosmotech_api.dataset.download.common import download_dataset_by_id
+from cosmotech.coal.utils.semver import semver_of
 
 
+@pytest.mark.skipif(semver_of("cosmotech_api").major >= 5, reason="not supported in version 5")
 class TestCommonFunctions:
     """Tests for top-level functions in the common module."""