Fix formatting and improve parsing of risk factor's parquet file (#48)

Author: read-write-content[bot]

armis_sdk/__init__.py

@@ -1,2 +1,2 @@
-from armis_sdk.core.armis_sdk import ArmisSdk
-from armis_sdk.core.client_credentials import ClientCredentials
+from armis_sdk.core.armis_sdk import ArmisSdk  # noqa: F401
+from armis_sdk.core.client_credentials import ClientCredentials  # noqa: F401

armis_sdk/clients/sites_client.py

@@ -1,5 +1,4 @@
 from typing import AsyncIterator
-from typing import List
 
 import universalasync
 
@@ -129,7 +128,7 @@ async def main():
             data = response_utils.get_data_dict(response)
             return Site.model_validate(data)
 
-    async def hierarchy(self) -> List[Site]:
+    async def hierarchy(self) -> list[Site]:
         """Create a hierarchy of the tenant's sites, taking into account the parent-child relationships.
 
         Returns:

armis_sdk/core/armis_error.py

@@ -4,7 +4,6 @@
 """
 
 import json
-from typing import List
 from typing import Optional
 from typing import Union
 
@@ -26,7 +25,7 @@ def __str__(self):
 
 
 class ErrorBody(BaseModel):
-    detail: Union[str, List[DetailItem]]
+    detail: Union[str, list[DetailItem]]
 
 
 class ArmisError(Exception):
@@ -63,7 +62,7 @@ class ResponseError(ArmisError):
     def __init__(
         self,
         error_body: ErrorBody,
-        response_errors: Optional[List[HTTPStatusError]] = None,
+        response_errors: Optional[list[HTTPStatusError]] = None,
     ):
         super().__init__(self._get_message(error_body))
         self.response_errors = response_errors

armis_sdk/entities/asq_rule.py

@@ -1,4 +1,3 @@
-from typing import List
 from typing import Optional
 from typing import Union
 
@@ -41,10 +40,10 @@ class AsqRule(BaseEntity):
         ```
     """
 
-    and_: Optional[List[Union[str, "AsqRule"]]] = Field(alias="and", default=None)
+    and_: Optional[list[Union[str, "AsqRule"]]] = Field(alias="and", default=None)
     """Rules that all must match."""
 
-    or_: Optional[List[Union[str, "AsqRule"]]] = Field(alias="or", default=None)
+    or_: Optional[list[Union[str, "AsqRule"]]] = Field(alias="or", default=None)
     """Rules that at least one of them must match."""
 
     @classmethod

armis_sdk/entities/data_export/risk_factor.py

@@ -68,7 +68,7 @@ class RiskFactor(BaseExportedEntity):
     **Example**: `Device Supports SMBv1`
     """
 
-    score: int
+    score: Optional[int]
     """The score of the risk factor"""
 
     group: str
@@ -78,14 +78,14 @@ class RiskFactor(BaseExportedEntity):
     **Example**: `INSECURE_TRAFFIC_AND_BEHAVIOR`
     """
 
-    remediation_type: str
+    remediation_type: Optional[str]
     """
     The type of the remediation
 
     **Example**: `Disable SMBv1 Protocol`
     """
 
-    remediation_description: str
+    remediation_description: Optional[str]
     """
     The description of the remediation
 
@@ -130,15 +130,27 @@ def series_to_model(cls, series: pandas.Series) -> "RiskFactor":
             category=series.loc["category"],
             type=series.loc["type"],
             description=series.loc["description"],
-            score=series.loc["score"],
+            score=(
+                int(score)
+                if (score := cls._value_or_none(series.loc["score"]))
+                else None
+            ),
             status=series.loc["status"],
             group=series.loc["group"],
-            remediation_type=series.loc["remediation"],
-            remediation_description=series.loc["remediation_description"],
-            remediation_recommended_actions=[
-                RiskFactorRecommendedAction(**item)
-                for item in json.loads(series.loc["remediation_recommended_actions"])
-            ],
+            remediation_type=cls._value_or_none(series.loc["remediation"]),
+            remediation_description=cls._value_or_none(
+                series.loc["remediation_description"]
+            ),
+            remediation_recommended_actions=(
+                [
+                    RiskFactorRecommendedAction(**item)
+                    for item in json.loads(
+                        series.loc["remediation_recommended_actions"]
+                    )
+                ]
+                if series.loc["remediation_recommended_actions"]
+                else []
+            ),
             first_seen=series.loc["first_seen"].to_pydatetime(),
             last_seen=series.loc["last_seen"].to_pydatetime(),
             status_update_time=cls._value_or_none(series.loc["status_update_time"]),

armis_sdk/entities/site.py

@@ -1,7 +1,6 @@
 import json
 from typing import Annotated
 from typing import Any
-from typing import List
 from typing import Optional
 
 from pydantic import BeforeValidator
@@ -71,16 +70,16 @@ class Site(BaseEntity):
     """The ASQ rule of the site."""
 
     network_equipment_device_ids: Annotated[
-        Optional[List[int]], BeforeValidator(ensure_list_of_ints)
+        Optional[list[int]], BeforeValidator(ensure_list_of_ints)
     ] = None
     """The ids of network equipment devices associated with the site."""
 
     integration_ids: Annotated[
-        Optional[List[int]], BeforeValidator(ensure_list_of_ints)
+        Optional[list[int]], BeforeValidator(ensure_list_of_ints)
     ] = None
     """The ids of the integration associated with the site."""
 
-    children: List["Site"] = Field(default_factory=list)
+    children: list["Site"] = Field(default_factory=list)
     """The sub-sites that are directly under this site 
     (their `parent_id` will match this site's `id`)."""
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "armis_sdk"
-version = "1.1.1"
+version = "1.1.2"
 description = "The Armis SDK is a package that encapsulates common use-cases for interacting with the Armis platform."
 authors = [
     { name = "Shai Lachmanovich", email = "[email protected]" },

tests/armis_sdk/clients/assets_client_test.py

@@ -374,10 +374,10 @@ async def test_update_with_failed_requests(httpx_mock: pytest_httpx.HTTPXMock):
     with pytest.raises(
         BulkUpdateError,
         match=(
-            "Failed to update item at index 2. "
-            'Request: {"asset_id": 2, "key": "custom.MyField1", '
-            '"operation": "SET", "value": "value3"}, '
-            'Response: {"status": 400, "error": "Bad Request"}'
+            r"Failed to update item at index 2\. "
+            r'Request: \{"asset_id": 2, "key": "custom\.MyField1", '
+            r'"operation": "SET", "value": "value3"\}, '
+            r'Response: \{"status": 400, "error": "Bad Request"\}'
         ),
     ):
         await assets_client.update(assets, fields)

tests/plugins/auto_setup_plugin.py

@@ -5,4 +5,4 @@
 
 @pytest.fixture(autouse=True)
 def auto_setup(setup_env_variables, authorized):  # pylint: disable=unused-argument
-    yield
+    return