diff --git a/README.md b/README.md index a7cf46a..2e87557 100644 --- a/README.md +++ b/README.md @@ -128,8 +128,7 @@ entity = j1.create_entity( entity_key='my-unique-key', entity_type='my_type', entity_class='MyClass', - properties=properties, - timestamp=int(time.time()) * 1000 # Optional, defaults to current datetime + properties=properties ) print(entity['entity']) diff --git a/examples/02_entity_management.py b/examples/02_entity_management.py index a618f0a..c415934 100644 --- a/examples/02_entity_management.py +++ b/examples/02_entity_management.py @@ -76,13 +76,10 @@ def create_entity_examples(j1): 'backupRetentionPeriod': 7, 'tag.Environment': 'production', 'tag.Team': 'data', - 'metadata': { - 'createdBy': 'terraform', - 'lastBackup': '2024-01-01T00:00:00Z', - 'maintenanceWindow': 'sun:03:00-sun:04:00' - } - }, - timestamp=int(time.time()) * 1000 + 'createdBy': 'terraform', + 'lastBackup': '2024-01-01T00:00:00Z', + 'maintenanceWindow': 'sun:03:00-sun:04:00' + } ) print(f"Created complex entity: {complex_entity['entity']['_id']}\n") @@ -122,16 +119,12 @@ def update_entity_examples(j1, entity_id): entity_id=entity_id, properties={ 'isActive': False, - 'maintenanceWindow': { - 'start': '2024-01-01T00:00:00Z', - 'end': '2024-01-01T04:00:00Z', - 'reason': 'scheduled_maintenance' - }, - 'metadata': { - 'maintenancePerformedBy': 'admin@company.com', - 'maintenanceType': 'security_patches', - 'estimatedDuration': '4 hours' - } + 'maintenanceWindowStart': '2024-01-01T00:00:00Z', + 'maintenanceWindowEnd': '2024-01-01T04:00:00Z', + 'maintenanceReason': 'scheduled_maintenance', + 'maintenancePerformedBy': 'admin@company.com', + 'maintenanceType': 'security_patches', + 'estimatedDuration': '4 hours' } ) print(f"Updated with complex properties\n") diff --git a/examples/03_relationship_management.py b/examples/03_relationship_management.py index 73037cb..04e7997 100644 --- a/examples/03_relationship_management.py +++ b/examples/03_relationship_management.py @@ -100,11 +100,9 @@ def create_relationship_examples(j1, from_entity_id, to_entity_id): 'version': '2.1.0', 'installPath': '/usr/local/bin/software', 'permissions': ['read', 'execute'], - 'metadata': { - 'installer': 'package_manager', - 'verified': True, - 'checksum': 'sha256:abc123...' - }, + 'installer': 'package_manager', + 'verified': True, + 'checksum': 'sha256:abc123...', 'tag.InstallationType': 'automated', 'tag.Verified': 'true' } @@ -142,11 +140,9 @@ def update_relationship_examples(j1, relationship_id, from_entity_id, to_entity_ 'lastModified': int(time.time()) * 1000, 'modifiedBy': 'security_team', 'expiresOn': int(time.time() + 86400) * 1000, # 24 hours from now - 'auditLog': { - 'previousLevel': 'write', - 'reason': 'promotion_requested', - 'approvedBy': 'security_manager' - } + 'previousLevel': 'write', + 'promotionReason': 'promotion_requested', + 'approvedBy': 'security_manager' } ) print(f"Updated with complex properties\n") diff --git a/examples/09_custom_file_transfer_example.py b/examples/09_custom_file_transfer_example.py index 67c75cd..5bbfacf 100644 --- a/examples/09_custom_file_transfer_example.py +++ b/examples/09_custom_file_transfer_example.py @@ -20,7 +20,11 @@ import sys # Add the parent directory to the path so we can import the jupiterone client -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +try: + sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +except NameError: + # Handle case when __file__ is not available (e.g., when exec'd) + sys.path.append('..') from jupiterone.client import JupiterOneClient diff --git a/examples/J1QLdeferredResponse.py b/examples/J1QLdeferredResponse.py index 244cf3a..8c69709 100644 --- a/examples/J1QLdeferredResponse.py +++ b/examples/J1QLdeferredResponse.py @@ -13,8 +13,8 @@ # JupiterOne GraphQL API headers j1_graphql_headers = { 'Content-Type': 'application/json', - 'Authorization': 'Bearer ' + token, - 'Jupiterone-Account': acct + 'Authorization': 'Bearer ' + (token or ''), + 'Jupiterone-Account': acct or '' } gql_query = """ diff --git a/examples/bulk_upload.py b/examples/bulk_upload.py index a8ac5ea..b64f68d 100644 --- a/examples/bulk_upload.py +++ b/examples/bulk_upload.py @@ -5,9 +5,15 @@ token = os.environ.get("JUPITERONE_TOKEN") url = "https://graphql.us.jupiterone.io" +# Check if credentials are available +if not account or not token: + print("Error: JUPITERONE_ACCOUNT and JUPITERONE_TOKEN environment variables must be set") + print("This example script requires valid JupiterOne credentials to run") + exit(1) + j1 = JupiterOneClient(account=account, token=token, url=url) -instance_id = "e7113c37-1ea8-4d00-9b82-c24952e70916" +instance_id = "" sync_job = j1.start_sync_job( instance_id=instance_id, diff --git a/examples/examples.py b/examples/examples.py index 061e509..6a710f0 100644 --- a/examples/examples.py +++ b/examples/examples.py @@ -8,6 +8,12 @@ token = os.environ.get("JUPITERONE_TOKEN") url = "https://graphql.us.jupiterone.io" +# Check if credentials are available +if not account or not token: + print("Error: JUPITERONE_ACCOUNT and JUPITERONE_TOKEN environment variables must be set") + print("This example script requires valid JupiterOne credentials to run") + exit(1) + j1 = JupiterOneClient(account=account, token=token, url=url) # query_v1 @@ -31,8 +37,7 @@ entity_key='jupiterone-api-client-python:{}'.format(num1), entity_type='python_client_create_entity', entity_class='Record', - properties=properties, - timestamp=int(time.time()) * 1000 # Optional, defaults to current datetime + properties=properties ) print("create_entity()") print(create_r) @@ -63,8 +68,7 @@ entity_key='jupiterone-api-client-python:{}'.format(num2), entity_type='python_client_create_entity', entity_class='Record', - properties=properties, - timestamp=int(time.time()) * 1000 # Optional, defaults to current datetime + properties=properties ) print("create_entity()") print(create_r_2) diff --git a/jupiterone/__init__.py b/jupiterone/__init__.py index 0d227bb..4199fb5 100644 --- a/jupiterone/__init__.py +++ b/jupiterone/__init__.py @@ -1,5 +1,13 @@ from .client import JupiterOneClient from .errors import ( JupiterOneClientError, - JupiterOneApiError + JupiterOneApiError, + JupiterOneApiRetryError ) + +__all__ = [ + "JupiterOneClient", + "JupiterOneClientError", + "JupiterOneApiError", + "JupiterOneApiRetryError" +] diff --git a/jupiterone/client.py b/jupiterone/client.py index ec4fa52..e5f9ae9 100644 --- a/jupiterone/client.py +++ b/jupiterone/client.py @@ -2,11 +2,12 @@ import json import os from warnings import warn -from typing import Dict, List, Union, Optional +from typing import Dict, List, Union, Optional, Any from datetime import datetime import time import re import requests +import urllib.parse from requests.adapters import HTTPAdapter, Retry import concurrent.futures @@ -67,28 +68,30 @@ class JupiterOneClient: # pylint: disable=too-many-instance-attributes - DEFAULT_URL = "https://graphql.us.jupiterone.io" - SYNC_API_URL = "https://api.us.jupiterone.io" + DEFAULT_URL: str = "https://graphql.us.jupiterone.io" + SYNC_API_URL: str = "https://api.us.jupiterone.io" def __init__( self, - account: str = None, - token: str = None, + account: Optional[str] = None, + token: Optional[str] = None, url: str = DEFAULT_URL, sync_url: str = SYNC_API_URL, - ): - self.account = account - self.token = token - self.graphql_url = url - self.sync_url = sync_url - self.headers = { - "Authorization": "Bearer {}".format(self.token), - "JupiterOne-Account": self.account, + ) -> None: + # Validate inputs + self._validate_constructor_inputs(account, token, url, sync_url) + self.account: Optional[str] = account + self.token: Optional[str] = token + self.graphql_url: str = url + self.sync_url: str = sync_url + self.headers: Dict[str, str] = { + "Authorization": "Bearer {}".format(self.token or ""), + "JupiterOne-Account": self.account or "", "Content-Type": "application/json", } # Initialize session with retry logic - self.session = requests.Session() + self.session: requests.Session = requests.Session() retries = Retry( total=5, backoff_factor=1, @@ -97,40 +100,134 @@ def __init__( ) self.session.mount("https://", HTTPAdapter(max_retries=retries)) + def _validate_constructor_inputs( + self, + account: Optional[str], + token: Optional[str], + url: str, + sync_url: str + ) -> None: + """Validate constructor inputs""" + # Validate account + if account is not None: + if not isinstance(account, str): + raise JupiterOneClientError("Account must be a string") + if not account.strip(): + raise JupiterOneClientError("Account cannot be empty") + if len(account) < 3: + raise JupiterOneClientError("Account ID appears to be too short") + + # Validate token + if token is not None: + if not isinstance(token, str): + raise JupiterOneClientError("Token must be a string") + if not token.strip(): + raise JupiterOneClientError("Token cannot be empty") + if len(token) < 10: + raise JupiterOneClientError("Token appears to be too short") + + # Validate URLs + self._validate_url(url, "GraphQL URL") + self._validate_url(sync_url, "Sync API URL") + + def _validate_url(self, url: str, url_name: str) -> None: + """Validate URL format""" + if not isinstance(url, str): + raise JupiterOneClientError(f"{url_name} must be a string") + + if not url.strip(): + raise JupiterOneClientError(f"{url_name} cannot be empty") + + try: + parsed = urllib.parse.urlparse(url) + if not parsed.scheme or not parsed.netloc: + raise ValueError("Invalid URL format") + if parsed.scheme not in ['http', 'https']: + raise ValueError("URL must use http or https protocol") + except Exception as e: + raise JupiterOneClientError(f"Invalid {url_name}: {str(e)}") + + def _validate_entity_id(self, entity_id: str, param_name: str = "entity_id") -> None: + """Validate entity ID format""" + if not isinstance(entity_id, str): + raise JupiterOneClientError(f"{param_name} must be a string") + + if not entity_id.strip(): + raise JupiterOneClientError(f"{param_name} cannot be empty") + + if len(entity_id) < 10: + raise JupiterOneClientError(f"{param_name} appears to be too short") + + def _validate_query_string(self, query: str, param_name: str = "query") -> None: + """Validate J1QL query string""" + if not isinstance(query, str): + raise JupiterOneClientError(f"{param_name} must be a string") + + if not query.strip(): + raise JupiterOneClientError(f"{param_name} cannot be empty") + + # Basic J1QL validation + query_upper = query.upper().strip() + if not query_upper.startswith('FIND'): + raise JupiterOneClientError(f"{param_name} must be a valid J1QL query starting with FIND (case-insensitive)") + + def _validate_properties(self, properties: Dict[str, Any], param_name: str = "properties") -> None: + """Validate entity/relationship properties""" + if not isinstance(properties, dict): + raise JupiterOneClientError(f"{param_name} must be a dictionary") + + # Check for nested objects (not supported by JupiterOne API) + for key, value in properties.items(): + if isinstance(value, dict): + raise JupiterOneClientError( + f"Nested objects in {param_name} are not supported by JupiterOne API. " + f"Key '{key}' contains a nested dictionary. Please flatten the structure." + ) + if isinstance(value, list) and any(isinstance(item, dict) for item in value): + raise JupiterOneClientError( + f"Lists containing dictionaries in {param_name} are not supported by JupiterOne API. " + f"Key '{key}' contains a list with dictionaries. Please flatten the structure." + ) + @property - def account(self): + def account(self) -> Optional[str]: """Your JupiterOne account ID""" return self._account @account.setter - def account(self, value: str): + def account(self, value: Optional[str]) -> None: """Your JupiterOne account ID""" if not value: raise JupiterOneClientError("account is required") self._account = value @property - def token(self): + def token(self) -> Optional[str]: """Your JupiterOne access token""" return self._token @token.setter - def token(self, value: str): + def token(self, value: Optional[str]) -> None: """Your JupiterOne access token""" if not value: raise JupiterOneClientError("token is required") self._token = value # pylint: disable=R1710 - def _execute_query(self, query: str, variables: Dict = None) -> Dict: + def _execute_query(self, query: str, variables: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """Executes query against graphql endpoint""" + # Validate credentials before making API calls + if not self.account: + raise JupiterOneClientError("Account is required. Please set the account property.") + if not self.token: + raise JupiterOneClientError("Token is required. Please set the token property.") - data = {"query": query} + data: Dict[str, Any] = {"query": query} if variables: - data.update(variables=variables) + data["variables"] = variables # Always ask for variableResultSize - data.update(flags={"variableResultSize": True}) + data["flags"] = {"variableResultSize": True} response = self.session.post( self.graphql_url, @@ -179,10 +276,10 @@ def _execute_query(self, query: str, variables: Dict = None) -> Dict: def _cursor_query( self, query: str, - cursor: str = None, + cursor: Optional[str] = None, include_deleted: bool = False, max_workers: Optional[int] = None - ) -> Dict: + ) -> Dict[str, Any]: """Performs a V1 graph query using cursor pagination with optional parallel processing args: @@ -201,9 +298,9 @@ def _cursor_query( else: result_limit = False - results: List = [] + results: List[Dict[str, Any]] = [] - def fetch_page(cursor: Optional[str] = None) -> Dict: + def fetch_page(cursor: Optional[str] = None) -> Dict[str, Any]: variables = {"query": query, "includeDeleted": include_deleted} if cursor is not None: variables["cursor"] = cursor @@ -287,8 +384,8 @@ def _limit_and_skip_query( skip: int = J1QL_SKIP_COUNT, limit: int = J1QL_LIMIT_COUNT, include_deleted: bool = False, - ) -> Dict: - results: List = [] + ) -> Dict[str, Any]: + results: List[Dict[str, Any]] = [] page: int = 0 while True: @@ -304,7 +401,7 @@ def _limit_and_skip_query( if "vertices" in data and "edges" in data: return data - if len(data) < J1QL_SKIP_COUNT: + if len(data) < skip: results.extend(data) break @@ -313,7 +410,7 @@ def _limit_and_skip_query( return {"data": results} - def query_with_deferred_response(self, query, cursor=None): + def query_with_deferred_response(self, query: str, cursor: Optional[str] = None) -> List[Dict[str, Any]]: """ Execute a J1QL query that returns a deferred response for handling large result sets. @@ -389,8 +486,13 @@ def query_with_deferred_response(self, query, cursor=None): return all_query_results - def _execute_syncapi_request(self, endpoint: str, payload: Dict = None) -> Dict: + def _execute_syncapi_request(self, endpoint: str, payload: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """Executes POST request to SyncAPI endpoints""" + # Validate credentials before making API calls + if not self.account: + raise JupiterOneClientError("Account is required. Please set the account property.") + if not self.token: + raise JupiterOneClientError("Token is required. Please set the token property.") # initiate requests session and implement retry logic of 5 request retries with 1 second between retries response = self.session.post( @@ -412,6 +514,7 @@ def _execute_syncapi_request(self, endpoint: str, payload: Dict = None) -> Dict: ) raise JupiterOneApiError(content.get("errors")) return response.json() + return {} elif response.status_code == 401: raise JupiterOneApiError( @@ -436,7 +539,7 @@ def _execute_syncapi_request(self, endpoint: str, payload: Dict = None) -> Dict: content = data.get("error", data.get("errors", content)) raise JupiterOneApiError("{}:{}".format(response.status_code, content)) - def query_v1(self, query: str, **kwargs) -> Dict: + def query_v1(self, query: str, **kwargs: Any) -> Dict[str, Any]: """Performs a V1 graph query args: query (str): Query text @@ -445,6 +548,25 @@ def query_v1(self, query: str, **kwargs) -> Dict: cursor (str): A pagination cursor for the initial query include_deleted (bool): Include recently deleted entities in query/search """ + # Validate inputs + self._validate_query_string(query) + + # Validate kwargs + if 'skip' in kwargs and kwargs['skip'] is not None: + if not isinstance(kwargs['skip'], int) or kwargs['skip'] < 0: + raise JupiterOneClientError("skip must be a non-negative integer") + + if 'limit' in kwargs and kwargs['limit'] is not None: + if not isinstance(kwargs['limit'], int) or kwargs['limit'] <= 0: + raise JupiterOneClientError("limit must be a positive integer") + + if 'cursor' in kwargs and kwargs['cursor'] is not None: + if not isinstance(kwargs['cursor'], str) or not kwargs['cursor'].strip(): + raise JupiterOneClientError("cursor must be a non-empty string") + + if 'include_deleted' in kwargs and kwargs['include_deleted'] is not None: + if not isinstance(kwargs['include_deleted'], bool): + raise JupiterOneClientError("include_deleted must be a boolean") uses_limit_and_skip: bool = "skip" in kwargs.keys() or "limit" in kwargs.keys() skip: int = kwargs.pop("skip", J1QL_SKIP_COUNT) limit: int = kwargs.pop("limit", J1QL_LIMIT_COUNT) @@ -467,34 +589,54 @@ def query_v1(self, query: str, **kwargs) -> Dict: query=query, cursor=cursor, include_deleted=include_deleted ) - def create_entity(self, **kwargs) -> Dict: + def create_entity(self, **kwargs: Any) -> Dict[str, Any]: """Creates an entity in graph. It will also update an existing entity. args: entity_key (str): Unique key for the entity entity_type (str): Value for _type of entity entity_class (str): Value for _class of entity - timestamp (int): Specify createdOn timestamp properties (dict): Dictionary of key/value entity properties """ + # Validate required parameters + entity_key = kwargs.get("entity_key") + entity_type = kwargs.get("entity_type") + entity_class = kwargs.get("entity_class") + + if not entity_key: + raise JupiterOneClientError("entity_key is required") + if not isinstance(entity_key, str) or not entity_key.strip(): + raise JupiterOneClientError("entity_key must be a non-empty string") + + if not entity_type: + raise JupiterOneClientError("entity_type is required") + if not isinstance(entity_type, str) or not entity_type.strip(): + raise JupiterOneClientError("entity_type must be a non-empty string") + + if not entity_class: + raise JupiterOneClientError("entity_class is required") + if not isinstance(entity_class, str) or not entity_class.strip(): + raise JupiterOneClientError("entity_class must be a non-empty string") + + # Validate properties if provided + if "properties" in kwargs and kwargs["properties"] is not None: + self._validate_properties(kwargs["properties"]) + variables = { "entityKey": kwargs.pop("entity_key"), "entityType": kwargs.pop("entity_type"), "entityClass": kwargs.pop("entity_class"), } - timestamp: int = kwargs.pop("timestamp", None) properties: Dict = kwargs.pop("properties", None) - if timestamp: - variables.update(timestamp=timestamp) if properties: variables.update(properties=properties) response = self._execute_query(query=CREATE_ENTITY, variables=variables) return response["data"]["createEntity"] - def delete_entity(self, entity_id: str = None, timestamp: int = None, hard_delete: bool = True) -> Dict: + def delete_entity(self, entity_id: Optional[str] = None, timestamp: Optional[int] = None, hard_delete: bool = True) -> Dict[str, Any]: """Deletes an entity from the graph. args: @@ -502,13 +644,26 @@ def delete_entity(self, entity_id: str = None, timestamp: int = None, hard_delet timestamp (int, optional): Timestamp for the deletion. Defaults to None. hard_delete (bool): Whether to perform a hard delete. Defaults to True. """ - variables = {"entityId": entity_id, "hardDelete": hard_delete} + # Validate required parameters + if not entity_id: + raise JupiterOneClientError("entity_id is required") + self._validate_entity_id(entity_id) + + # Validate timestamp if provided + if timestamp is not None: + if not isinstance(timestamp, int) or timestamp <= 0: + raise JupiterOneClientError("timestamp must be a positive integer") + + # Validate hard_delete + if not isinstance(hard_delete, bool): + raise JupiterOneClientError("hard_delete must be a boolean") + variables: Dict[str, Any] = {"entityId": entity_id, "hardDelete": hard_delete} if timestamp: variables["timestamp"] = timestamp response = self._execute_query(DELETE_ENTITY, variables=variables) return response["data"]["deleteEntityV2"] - def update_entity(self, entity_id: str = None, properties: Dict = None) -> Dict: + def update_entity(self, entity_id: Optional[str] = None, properties: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """ Update an existing entity. @@ -516,11 +671,19 @@ def update_entity(self, entity_id: str = None, properties: Dict = None) -> Dict: entity_id (str): The _id of the entity to update properties (dict): Dictionary of key/value entity properties """ + # Validate required parameters + if not entity_id: + raise JupiterOneClientError("entity_id is required") + self._validate_entity_id(entity_id) + + if not properties: + raise JupiterOneClientError("properties is required") + self._validate_properties(properties) variables = {"entityId": entity_id, "properties": properties} response = self._execute_query(UPDATE_ENTITY, variables=variables) return response["data"]["updateEntity"] - def create_relationship(self, **kwargs) -> Dict: + def create_relationship(self, **kwargs: Any) -> Dict[str, Any]: """ Create a relationship (edge) between two entities (vertices). @@ -531,6 +694,39 @@ def create_relationship(self, **kwargs) -> Dict: from_entity_id (str): Entity ID of the source vertex to_entity_id (str): Entity ID of the destination vertex """ + # Validate required parameters + relationship_key = kwargs.get("relationship_key") + relationship_type = kwargs.get("relationship_type") + relationship_class = kwargs.get("relationship_class") + from_entity_id = kwargs.get("from_entity_id") + to_entity_id = kwargs.get("to_entity_id") + + if not relationship_key: + raise JupiterOneClientError("relationship_key is required") + if not isinstance(relationship_key, str) or not relationship_key.strip(): + raise JupiterOneClientError("relationship_key must be a non-empty string") + + if not relationship_type: + raise JupiterOneClientError("relationship_type is required") + if not isinstance(relationship_type, str) or not relationship_type.strip(): + raise JupiterOneClientError("relationship_type must be a non-empty string") + + if not relationship_class: + raise JupiterOneClientError("relationship_class is required") + if not isinstance(relationship_class, str) or not relationship_class.strip(): + raise JupiterOneClientError("relationship_class must be a non-empty string") + + if not from_entity_id: + raise JupiterOneClientError("from_entity_id is required") + self._validate_entity_id(from_entity_id, "from_entity_id") + + if not to_entity_id: + raise JupiterOneClientError("to_entity_id is required") + self._validate_entity_id(to_entity_id, "to_entity_id") + + # Validate properties if provided + if "properties" in kwargs and kwargs["properties"] is not None: + self._validate_properties(kwargs["properties"]) variables = { "relationshipKey": kwargs.pop("relationship_key"), "relationshipType": kwargs.pop("relationship_type"), @@ -546,7 +742,7 @@ def create_relationship(self, **kwargs) -> Dict: response = self._execute_query(query=CREATE_RELATIONSHIP, variables=variables) return response["data"]["createRelationship"] - def update_relationship(self, **kwargs) -> Dict: + def update_relationship(self, **kwargs: Any) -> Dict[str, Any]: """ Update a relationship (edge) between two entities (vertices). @@ -568,7 +764,7 @@ def update_relationship(self, **kwargs) -> Dict: response = self._execute_query(query=UPDATE_RELATIONSHIP, variables=variables) return response["data"]["updateRelationship"] - def delete_relationship(self, relationship_id: str = None): + def delete_relationship(self, relationship_id: Optional[str] = None) -> Dict[str, Any]: """Deletes a relationship between two entities. args: @@ -581,11 +777,11 @@ def delete_relationship(self, relationship_id: str = None): def create_integration_instance( self, - instance_name: str = None, - instance_description: str = None, + instance_name: Optional[str] = None, + instance_description: Optional[str] = None, integration_definition_id: str = "8013680b-311a-4c2e-b53b-c8735fd97a5c", - resource_group_id: str = None, - ): + resource_group_id: Optional[str] = None, + ) -> Dict[str, Any]: """Creates a new Custom Integration Instance. args: @@ -614,7 +810,7 @@ def create_integration_instance( response = self._execute_query(CREATE_INSTANCE, variables=variables) return response["data"]["createIntegrationInstance"] - def fetch_all_entity_properties(self): + def fetch_all_entity_properties(self) -> List[Dict[str, Any]]: """Fetch list of aggregated property keys from all entities in the graph.""" response = self._execute_query(query=ALL_PROPERTIES) @@ -629,7 +825,7 @@ def fetch_all_entity_properties(self): return return_list - def fetch_all_entity_tags(self): + def fetch_all_entity_tags(self) -> List[Dict[str, Any]]: """Fetch list of aggregated property keys from all entities in the graph.""" response = self._execute_query(query=ALL_PROPERTIES) @@ -644,7 +840,7 @@ def fetch_all_entity_tags(self): return return_list - def fetch_entity_raw_data(self, entity_id: str = None): + def fetch_entity_raw_data(self, entity_id: Optional[str] = None) -> Dict[str, Any]: """Fetch the contents of raw data for a given entity in a J1 Account.""" variables = {"entityId": entity_id, "source": "integration-managed"} @@ -654,10 +850,10 @@ def fetch_entity_raw_data(self, entity_id: str = None): def start_sync_job( self, - instance_id: str = None, - sync_mode: str = None, - source: str = None, - ): + instance_id: Optional[str] = None, + sync_mode: Optional[str] = None, + source: Optional[str] = None, + ) -> Dict[str, Any]: """Start a synchronization job. args: @@ -754,7 +950,7 @@ def bulk_delete_entities( return response - def finalize_sync_job(self, instance_job_id: str = None): + def finalize_sync_job(self, instance_job_id: Optional[str] = None) -> Dict[str, Any]: """Start a synchronization job. args: @@ -768,7 +964,7 @@ def finalize_sync_job(self, instance_job_id: str = None): return response - def fetch_integration_jobs(self, instance_id: str = None): + def fetch_integration_jobs(self, instance_id: Optional[str] = None) -> List[Dict[str, Any]]: """Fetch Integration Job details from defined integration instance. args: @@ -801,21 +997,21 @@ def fetch_integration_job_events( return response["data"]["integrationEvents"] - def get_integration_definition_details(self, integration_type: str = None): + def get_integration_definition_details(self, integration_type: Optional[str] = None) -> Dict[str, Any]: """Fetch the Integration Definition Details for a given integration type.""" variables = {"integrationType": integration_type, "includeConfig": True} response = self._execute_query(FIND_INTEGRATION_DEFINITION, variables=variables) return response - def fetch_integration_instances(self, definition_id: str = None): + def fetch_integration_instances(self, definition_id: Optional[str] = None) -> List[Dict[str, Any]]: """Fetch all configured Instances for a given integration type.""" variables = {"definitionId": definition_id, "limit": 100} response = self._execute_query(INTEGRATION_INSTANCES, variables=variables) return response - def get_integration_instance_details(self, instance_id: str = None): + def get_integration_instance_details(self, instance_id: Optional[str] = None) -> Dict[str, Any]: """Fetch configuration details for a single configured Integration Instance.""" variables = {"integrationInstanceId": instance_id} @@ -964,7 +1160,7 @@ def generate_j1ql(self, natural_language_prompt: str = None): return response["data"]["j1qlFromNaturalLanguage"] - def list_alert_rules(self): + def list_alert_rules(self) -> List[Dict[str, Any]]: """List all defined Alert Rules configured in J1 account""" results = [] @@ -993,7 +1189,7 @@ def list_alert_rules(self): return results - def get_alert_rule_details(self, rule_id: str = None): + def get_alert_rule_details(self, rule_id: Optional[str] = None) -> Dict[str, Any]: """Get details of a single defined Alert Rule configured in J1 account""" results = [] @@ -1100,7 +1296,7 @@ def create_alert_rule( return response["data"]["createInlineQuestionRuleInstance"] - def delete_alert_rule(self, rule_id: str = None): + def delete_alert_rule(self, rule_id: Optional[str] = None) -> Dict[str, Any]: """Delete a single Alert Rule configured in J1 account""" variables = { "id": rule_id @@ -1320,7 +1516,7 @@ def fetch_downloaded_evaluation_results(self, download_url: str = None): return e - def list_questions(self, search_query: str = None, tags: List[str] = None): + def list_questions(self, search_query: Optional[str] = None, tags: Optional[List[str]] = None) -> Dict[str, Any]: """List all defined Questions configured in J1 Account Questions Library Args: @@ -1392,7 +1588,7 @@ def list_questions(self, search_query: str = None, tags: List[str] = None): return results - def get_question_details(self, question_id: str = None): + def get_question_details(self, question_id: Optional[str] = None) -> Dict[str, Any]: """Get details of a specific question by ID Args: @@ -1422,9 +1618,9 @@ def get_question_details(self, question_id: str = None): def create_question( self, title: str, - queries: List[Dict], - **kwargs - ): + queries: List[Dict[str, Any]], + **kwargs: Any + ) -> Dict[str, Any]: """Creates a new Question in the J1 account. Args: @@ -1510,12 +1706,12 @@ def create_question( def update_question( self, question_id: str, - title: str = None, - description: str = None, - queries: List[Dict] = None, - tags: List[str] = None, - **kwargs - ) -> Dict: + title: Optional[str] = None, + description: Optional[str] = None, + queries: Optional[List[Dict[str, Any]]] = None, + tags: Optional[List[str]] = None, + **kwargs: Any + ) -> Dict[str, Any]: """ Update an existing question in the J1 account. @@ -1629,7 +1825,7 @@ def update_question( response = self._execute_query(UPDATE_QUESTION, variables) return response["data"]["updateQuestion"] - def delete_question(self, question_id: str) -> Dict: + def delete_question(self, question_id: str) -> Dict[str, Any]: """ Delete an existing question from the J1 account. @@ -1729,7 +1925,7 @@ def create_update_parameter( response = self._execute_query(UPSERT_PARAMETER, variables=variables) return response - def update_entity_v2(self, entity_id: str = None, properties: Dict = None) -> Dict: + def update_entity_v2(self, entity_id: Optional[str] = None, properties: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """ Update an existing entity by adding new or updating existing properties. @@ -1746,7 +1942,7 @@ def update_entity_v2(self, entity_id: str = None, properties: Dict = None) -> Di response = self._execute_query(UPDATE_ENTITYV2, variables=variables) return response["data"]["updateEntityV2"] - def get_cft_upload_url(self, integration_instance_id: str, filename: str, dataset_id: str) -> Dict: + def get_cft_upload_url(self, integration_instance_id: str, filename: str, dataset_id: str) -> Dict[str, Any]: """ Get an upload URL for Custom File Transfer integration. @@ -1792,7 +1988,7 @@ def get_cft_upload_url(self, integration_instance_id: str, filename: str, datase response = self._execute_query(query, variables) return response["data"]["integrationFileTransferUploadUrl"] - def upload_cft_file(self, upload_url: str, file_path: str) -> Dict: + def upload_cft_file(self, upload_url: str, file_path: str) -> Dict[str, Any]: """ Upload a CSV file to the Custom File Transfer integration using the provided upload URL. diff --git a/jupiterone/errors.py b/jupiterone/errors.py index eccf017..018a5ef 100644 --- a/jupiterone/errors.py +++ b/jupiterone/errors.py @@ -1,11 +1,14 @@ class JupiterOneClientError(Exception): - """ Raised when error creating client """ + """ Raised when error creating client """ + pass class JupiterOneApiRetryError(Exception): """ Used to trigger retry on rate limit """ + pass class JupiterOneApiError(Exception): """ Raised when API returns error response """ + pass diff --git a/setup.py b/setup.py index 57b1d2d..318fe06 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name="jupiterone", - version="2.0.0", + version="2.0.1", description="A Python client for the JupiterOne API", license="MIT License", author="JupiterOne", diff --git a/tests/test_alert_rule_methods.py b/tests/test_alert_rule_methods.py index 0634a40..f373afa 100644 --- a/tests/test_alert_rule_methods.py +++ b/tests/test_alert_rule_methods.py @@ -216,6 +216,10 @@ def test_update_alert_rule_basic(self, mock_execute_query, mock_get_details): "pollingInterval": "ONE_DAY", "tags": ["old-tag"], "labels": [], + "triggerActionsOnNewEntitiesOnly": True, + "ignorePreviousResults": False, + "notifyOnFailure": True, + "templates": {}, "operations": [{ "__typename": "Operation", "when": {"type": "FILTER", "condition": ["AND", ["queries.query0.total", ">", 0]]}, @@ -300,24 +304,25 @@ def test_fetch_evaluation_result_download_url(self, mock_execute_query): assert result == mock_response mock_execute_query.assert_called_once() - @patch('jupiterone.client.requests.get') - def test_fetch_downloaded_evaluation_results_success(self, mock_get): + def test_fetch_downloaded_evaluation_results_success(self): """Test fetch_downloaded_evaluation_results method - success""" mock_response = Mock() mock_response.json.return_value = {"data": [{"id": "result-1"}]} - mock_get.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.get.return_value = mock_response - result = self.client.fetch_downloaded_evaluation_results(download_url="https://example.com/download") + result = self.client.fetch_downloaded_evaluation_results(download_url="https://example.com/download") - assert result == {"data": [{"id": "result-1"}]} - mock_get.assert_called_once() + assert result == {"data": [{"id": "result-1"}]} + mock_session.get.assert_called_once_with("https://example.com/download", timeout=60) - @patch('jupiterone.client.requests.get') - def test_fetch_downloaded_evaluation_results_exception(self, mock_get): + def test_fetch_downloaded_evaluation_results_exception(self): """Test fetch_downloaded_evaluation_results method - exception""" - mock_get.side_effect = Exception("Network error") + with patch.object(self.client, 'session') as mock_session: + mock_session.get.side_effect = Exception("Network error") - result = self.client.fetch_downloaded_evaluation_results(download_url="https://example.com/download") + result = self.client.fetch_downloaded_evaluation_results(download_url="https://example.com/download") - assert isinstance(result, Exception) - assert str(result) == "Network error" \ No newline at end of file + assert isinstance(result, Exception) + assert str(result) == "Network error" \ No newline at end of file diff --git a/tests/test_client_init.py b/tests/test_client_init.py index e12d61b..39811e2 100644 --- a/tests/test_client_init.py +++ b/tests/test_client_init.py @@ -45,12 +45,12 @@ def test_client_init_missing_token(self): def test_client_init_empty_account(self): """Test client initialization with empty account""" - with pytest.raises(JupiterOneClientError, match="account is required"): + with pytest.raises(JupiterOneClientError, match="Account cannot be empty"): JupiterOneClient(account="", token="test-token") def test_client_init_empty_token(self): """Test client initialization with empty token""" - with pytest.raises(JupiterOneClientError, match="token is required"): + with pytest.raises(JupiterOneClientError, match="Token cannot be empty"): JupiterOneClient(account="test-account", token="") def test_client_init_none_account(self): @@ -127,127 +127,137 @@ def setup_method(self): """Set up test fixtures""" self.client = JupiterOneClient(account="test-account", token="test-token") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_401_error(self, mock_session): + def test_execute_query_401_error(self): """Test _execute_query method with 401 error""" mock_response = Mock() mock_response.status_code = 401 - mock_session.post.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response - with pytest.raises(JupiterOneApiError, match="401: Unauthorized"): - self.client._execute_query("test query") + with pytest.raises(JupiterOneApiError, match="401: Unauthorized"): + self.client._execute_query("test query") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_429_error(self, mock_session): + def test_execute_query_429_error(self): """Test _execute_query method with 429 error""" mock_response = Mock() mock_response.status_code = 429 - mock_session.post.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response - with pytest.raises(JupiterOneApiRetryError, match="rate limit exceeded"): - self.client._execute_query("test query") + with pytest.raises(JupiterOneApiRetryError, match="rate limit exceeded"): + self.client._execute_query("test query") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_503_error(self, mock_session): + def test_execute_query_503_error(self): """Test _execute_query method with 503 error""" mock_response = Mock() mock_response.status_code = 503 - mock_session.post.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response - with pytest.raises(JupiterOneApiRetryError, match="rate limit exceeded"): - self.client._execute_query("test query") + with pytest.raises(JupiterOneApiRetryError, match="rate limit exceeded"): + self.client._execute_query("test query") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_504_error(self, mock_session): + def test_execute_query_504_error(self): """Test _execute_query method with 504 error""" mock_response = Mock() mock_response.status_code = 504 - mock_session.post.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response - with pytest.raises(JupiterOneApiRetryError, match="Gateway Timeout"): - self.client._execute_query("test query") + with pytest.raises(JupiterOneApiRetryError, match="Gateway Timeout"): + self.client._execute_query("test query") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_500_error(self, mock_session): + def test_execute_query_500_error(self): """Test _execute_query method with 500 error""" mock_response = Mock() mock_response.status_code = 500 - mock_session.post.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response - with pytest.raises(JupiterOneApiError, match="internal server error"): - self.client._execute_query("test query") + with pytest.raises(JupiterOneApiError, match="internal server error"): + self.client._execute_query("test query") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_200_with_errors(self, mock_session): + def test_execute_query_200_with_errors(self): """Test _execute_query method with 200 status but GraphQL errors""" mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = { "errors": [{"message": "GraphQL error"}] } - mock_session.post.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response - with pytest.raises(JupiterOneApiError): - self.client._execute_query("test query") + with pytest.raises(JupiterOneApiError): + self.client._execute_query("test query") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_200_with_429_in_errors(self, mock_session): + def test_execute_query_200_with_429_in_errors(self): """Test _execute_query method with 200 status but 429 in GraphQL errors""" mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = { "errors": [{"message": "429 rate limit exceeded"}] } - mock_session.post.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response - with pytest.raises(JupiterOneApiRetryError, match="rate limit exceeded"): - self.client._execute_query("test query") + with pytest.raises(JupiterOneApiRetryError, match="rate limit exceeded"): + self.client._execute_query("test query") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_200_success(self, mock_session): + def test_execute_query_200_success(self): """Test _execute_query method with successful 200 response""" mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = { "data": {"result": "success"} } - mock_session.post.return_value = mock_response - - result = self.client._execute_query("test query") - assert result == {"data": {"result": "success"}} + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response + + result = self.client._execute_query("test query") + + assert result == {"data": {"result": "success"}} - @patch.object(JupiterOneClient, 'session') - def test_execute_query_with_variables(self, mock_session): + def test_execute_query_with_variables(self): """Test _execute_query method with variables""" mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = { "data": {"result": "success"} } - mock_session.post.return_value = mock_response - - variables = {"key": "value"} - self.client._execute_query("test query", variables=variables) - # Verify that variables were included in the request - call_args = mock_session.post.call_args - assert "variables" in call_args[1]["json"] - assert call_args[1]["json"]["variables"] == variables - - @patch.object(JupiterOneClient, 'session') - def test_execute_query_with_flags(self, mock_session): + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response + + variables = {"key": "value"} + self.client._execute_query("test query", variables=variables) + + # Verify that variables were included in the request + call_args = mock_session.post.call_args + assert "variables" in call_args[1]["json"] + assert call_args[1]["json"]["variables"] == variables + + def test_execute_query_with_flags(self): """Test _execute_query method includes flags""" mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = { "data": {"result": "success"} } - mock_session.post.return_value = mock_response - - self.client._execute_query("test query") - # Verify that flags were included in the request - call_args = mock_session.post.call_args - assert "flags" in call_args[1]["json"] - assert call_args[1]["json"]["flags"] == {"variableResultSize": True} \ No newline at end of file + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response + + self.client._execute_query("test query") + + # Verify that flags were included in the request + call_args = mock_session.post.call_args + assert "flags" in call_args[1]["json"] + assert call_args[1]["json"]["flags"] == {"variableResultSize": True} \ No newline at end of file diff --git a/tests/test_create_entity.py b/tests/test_create_entity.py index b31b316..cc18fdf 100644 --- a/tests/test_create_entity.py +++ b/tests/test_create_entity.py @@ -37,7 +37,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.create_entity( entity_key='host1', entity_type='test_host', diff --git a/tests/test_create_relationship.py b/tests/test_create_relationship.py index f63d4b5..c53bb65 100644 --- a/tests/test_create_relationship.py +++ b/tests/test_create_relationship.py @@ -41,13 +41,13 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.create_relationship( relationship_key='relationship1', relationship_type='test_relationship', relationship_class='TestRelationship', - from_entity_id='2', - to_entity_id='1' + from_entity_id='entity-id-12345', + to_entity_id='entity-id-67890' ) assert type(response) == dict diff --git a/tests/test_cursor_query_edge_cases.py b/tests/test_cursor_query_edge_cases.py index 8c8cb20..6aafe30 100644 --- a/tests/test_cursor_query_edge_cases.py +++ b/tests/test_cursor_query_edge_cases.py @@ -211,7 +211,7 @@ def test_cursor_query_with_include_deleted_true(self, mock_execute_query): # Verify the query was called with includeDeleted=True mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - variables = call_args[0][1] + variables = call_args[1]["variables"] assert variables["includeDeleted"] is True # Verify the result @@ -241,7 +241,7 @@ def test_cursor_query_with_include_deleted_false(self, mock_execute_query): # Verify the query was called with includeDeleted=False mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - variables = call_args[0][1] + variables = call_args[1]["variables"] assert variables["includeDeleted"] is False # Verify the result @@ -271,7 +271,7 @@ def test_cursor_query_with_initial_cursor(self, mock_execute_query): # Verify the query was called with the initial cursor mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - variables = call_args[0][1] + variables = call_args[1]["variables"] assert variables["cursor"] == "initial_cursor" # Verify the result diff --git a/tests/test_delete_entity.py b/tests/test_delete_entity.py index 16c123d..f232ddf 100644 --- a/tests/test_delete_entity.py +++ b/tests/test_delete_entity.py @@ -32,7 +32,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.delete_entity('1') assert type(response) == dict @@ -68,7 +68,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.delete_entity('2', timestamp=1640995200000) assert type(response) == dict @@ -103,7 +103,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.delete_entity('3', hard_delete=False) assert type(response) == dict @@ -138,7 +138,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.delete_entity('4', timestamp=1640995200000, hard_delete=True) assert type(response) == dict diff --git a/tests/test_delete_relationship.py b/tests/test_delete_relationship.py index 42493b3..dba6d5d 100644 --- a/tests/test_delete_relationship.py +++ b/tests/test_delete_relationship.py @@ -41,7 +41,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.delete_relationship('1') assert type(response) == dict diff --git a/tests/test_limit_skip_query_edge_cases.py b/tests/test_limit_skip_query_edge_cases.py index 7a83250..926cf6e 100644 --- a/tests/test_limit_skip_query_edge_cases.py +++ b/tests/test_limit_skip_query_edge_cases.py @@ -79,7 +79,9 @@ def test_limit_and_skip_query_multiple_pages_with_break(self, mock_execute_query "queryV1": { "data": [ {"id": "1", "name": "entity1"}, - {"id": "2", "name": "entity2"} + {"id": "2", "name": "entity2"}, + {"id": "3", "name": "entity3"}, + {"id": "4", "name": "entity4"} ] } } @@ -98,7 +100,7 @@ def test_limit_and_skip_query_multiple_pages_with_break(self, mock_execute_query mock_execute_query.side_effect = [mock_response1, mock_response2] - result = self.client._limit_and_skip_query("FIND * LIMIT 10") + result = self.client._limit_and_skip_query("FIND * LIMIT 10", skip=3) # Should call twice, but break on second page assert mock_execute_query.call_count == 2 @@ -107,6 +109,8 @@ def test_limit_and_skip_query_multiple_pages_with_break(self, mock_execute_query assert result == {"data": [ {"id": "1", "name": "entity1"}, {"id": "2", "name": "entity2"}, + {"id": "3", "name": "entity3"}, + {"id": "4", "name": "entity4"}, {"id": "3", "name": "entity3"} ]} @@ -134,8 +138,8 @@ def test_limit_and_skip_query_with_custom_skip_limit(self, mock_execute_query): # Verify the query was called with custom skip/limit mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - query = call_args[0][0] - assert "SKIP 0 LIMIT 25" in query # First page starts at 0 + variables = call_args[1]["variables"] + assert "SKIP 0 LIMIT 25" in variables["query"] # First page starts at 0 # Verify the result assert result == {"data": [ @@ -151,7 +155,9 @@ def test_limit_and_skip_query_multiple_pages_with_custom_values(self, mock_execu "queryV1": { "data": [ {"id": "1", "name": "entity1"}, - {"id": "2", "name": "entity2"} + {"id": "2", "name": "entity2"}, + {"id": "3", "name": "entity3"}, + {"id": "4", "name": "entity4"} ] } } @@ -172,7 +178,7 @@ def test_limit_and_skip_query_multiple_pages_with_custom_values(self, mock_execu result = self.client._limit_and_skip_query( "FIND * LIMIT 10", - skip=10, + skip=3, limit=10 ) @@ -181,13 +187,15 @@ def test_limit_and_skip_query_multiple_pages_with_custom_values(self, mock_execu # Verify the queries were called with correct skip values call_args_list = mock_execute_query.call_args_list - assert "SKIP 0 LIMIT 10" in call_args_list[0][0][0] # First page - assert "SKIP 10 LIMIT 10" in call_args_list[1][0][0] # Second page + assert "SKIP 0 LIMIT 10" in call_args_list[0][1]["variables"]["query"] # First page + assert "SKIP 3 LIMIT 10" in call_args_list[1][1]["variables"]["query"] # Second page # Verify the result assert result == {"data": [ {"id": "1", "name": "entity1"}, {"id": "2", "name": "entity2"}, + {"id": "3", "name": "entity3"}, + {"id": "4", "name": "entity4"}, {"id": "3", "name": "entity3"} ]} @@ -215,7 +223,7 @@ def test_limit_and_skip_query_with_include_deleted_true(self, mock_execute_query # Verify the query was called with includeDeleted=True mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - variables = call_args[0][1] + variables = call_args[1]["variables"] assert variables["includeDeleted"] is True # Verify the result @@ -248,7 +256,7 @@ def test_limit_and_skip_query_with_include_deleted_false(self, mock_execute_quer # Verify the query was called with includeDeleted=False mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - variables = call_args[0][1] + variables = call_args[1]["variables"] assert variables["includeDeleted"] is False # Verify the result @@ -305,11 +313,11 @@ def test_limit_and_skip_query_complex_query(self, mock_execute_query): # Verify the query was called with the complex query mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - query = call_args[0][0] - assert "FIND aws_instance" in query - assert "THAT RELATES TO aws_vpc" in query - assert "WITH tag.Environment = 'production'" in query - assert "SKIP 0 LIMIT" in query # Should have skip/limit added + variables = call_args[1]["variables"] + assert "FIND aws_instance" in variables["query"] + assert "THAT RELATES TO aws_vpc" in variables["query"] + assert "WITH tag.Environment = 'production'" in variables["query"] + assert "SKIP 0 LIMIT" in variables["query"] # Should have skip/limit added # Verify the result assert result == {"data": [ @@ -325,7 +333,9 @@ def test_limit_and_skip_query_pagination_math(self, mock_execute_query): "queryV1": { "data": [ {"id": "1", "name": "entity1"}, - {"id": "2", "name": "entity2"} + {"id": "2", "name": "entity2"}, + {"id": "3", "name": "entity3"}, + {"id": "4", "name": "entity4"} ] } } @@ -337,7 +347,9 @@ def test_limit_and_skip_query_pagination_math(self, mock_execute_query): "queryV1": { "data": [ {"id": "3", "name": "entity3"}, - {"id": "4", "name": "entity4"} + {"id": "4", "name": "entity4"}, + {"id": "5", "name": "entity5"}, + {"id": "6", "name": "entity6"} ] } } @@ -356,16 +368,16 @@ def test_limit_and_skip_query_pagination_math(self, mock_execute_query): mock_execute_query.side_effect = [mock_response1, mock_response2, mock_response3] - result = self.client._limit_and_skip_query("FIND * LIMIT 10") + result = self.client._limit_and_skip_query("FIND * LIMIT 10", skip=3) # Should call three times assert mock_execute_query.call_count == 3 # Verify the pagination math call_args_list = mock_execute_query.call_args_list - assert "SKIP 0 LIMIT" in call_args_list[0][0][0] # Page 0: SKIP 0 - assert "SKIP 100 LIMIT" in call_args_list[1][0][0] # Page 1: SKIP 100 - assert "SKIP 200 LIMIT" in call_args_list[2][0][0] # Page 2: SKIP 200 + assert "SKIP 0 LIMIT" in call_args_list[0][1]["variables"]["query"] # Page 0: SKIP 0 + assert "SKIP 3 LIMIT" in call_args_list[1][1]["variables"]["query"] # Page 1: SKIP 3 + assert "SKIP 6 LIMIT" in call_args_list[2][1]["variables"]["query"] # Page 2: SKIP 6 # Verify the result assert result == {"data": [ @@ -373,6 +385,10 @@ def test_limit_and_skip_query_pagination_math(self, mock_execute_query): {"id": "2", "name": "entity2"}, {"id": "3", "name": "entity3"}, {"id": "4", "name": "entity4"}, + {"id": "3", "name": "entity3"}, + {"id": "4", "name": "entity4"}, + {"id": "5", "name": "entity5"}, + {"id": "6", "name": "entity6"}, {"id": "5", "name": "entity5"} ]} @@ -396,8 +412,8 @@ def test_limit_and_skip_query_default_constants(self, mock_execute_query): # Verify the query was called with default constants mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - query = call_args[0][0] - assert f"SKIP 0 LIMIT {J1QL_LIMIT_COUNT}" in query + variables = call_args[1]["variables"] + assert f"SKIP 0 LIMIT {J1QL_LIMIT_COUNT}" in variables["query"] # Verify the result assert result == {"data": [ diff --git a/tests/test_query.py b/tests/test_query.py index 4ae28c4..890a37f 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -83,7 +83,7 @@ def test_execute_query(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" variables = { 'query': query, @@ -110,7 +110,7 @@ def test_limit_skip_query_v1(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" response = j1.query_v1( query=query, @@ -139,7 +139,7 @@ def test_cursor_query_v1(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" response = j1.query_v1( @@ -186,7 +186,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1' return tree" response = j1.query_v1( query=query, @@ -236,7 +236,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1' return tree" response = j1.query_v1(query) @@ -268,7 +268,7 @@ def test_retry_on_limit_skip_query(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" response = j1.query_v1( query=query, @@ -302,7 +302,7 @@ def test_retry_on_cursor_query(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" response = j1.query_v1( query=query @@ -322,7 +322,7 @@ def test_avoid_retry_on_limit_skip_query(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" with pytest.raises(JupiterOneApiError): j1.query_v1( @@ -340,7 +340,7 @@ def test_avoid_retry_on_cursor_query(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" with pytest.raises(JupiterOneApiError): j1.query_v1( @@ -356,7 +356,7 @@ def test_warn_limit_and_skip_deprecated(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" with pytest.warns(DeprecationWarning): diff --git a/tests/test_update_entity.py b/tests/test_update_entity.py index cbd24c3..37b0514 100644 --- a/tests/test_update_entity.py +++ b/tests/test_update_entity.py @@ -34,7 +34,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.update_entity('1', properties={'testKey': 'testValue'}) assert type(response) == dict