Custom User-Agent header, Migrated all bare requests.post() to self.session.post(), Version bump to 2.3.0

jupiterone/client.py

@@ -68,6 +68,7 @@ class JupiterOneClient:
 
     # pylint: disable=too-many-instance-attributes
 
+    SDK_VERSION: str = "2.3.0"
     DEFAULT_URL: str = "https://graphql.us.jupiterone.io"
     SYNC_API_URL: str = "https://api.us.jupiterone.io"
 
@@ -77,21 +78,28 @@ def __init__(
         token: Optional[str] = None,
         url: str = DEFAULT_URL,
         sync_url: str = SYNC_API_URL,
+        user_agent: Optional[str] = None,
     ) -> 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
+
+        base_ua = f"jupiterone-client-python/{self.SDK_VERSION}"
+        full_ua = f"{base_ua} {user_agent}" if user_agent else base_ua
+
         self.headers: Dict[str, str] = {
             "Authorization": "Bearer {}".format(self.token or ""),
             "JupiterOne-Account": self.account or "",
             "Content-Type": "application/json",
+            "User-Agent": full_ua,
         }
 
         # Initialize session with retry logic
         self.session: requests.Session = requests.Session()
+        self.session.headers["User-Agent"] = full_ua
         retries = Retry(
             total=5,
             backoff_factor=1,
@@ -1204,8 +1212,8 @@ def list_alert_rules(self) -> List[Dict[str, Any]]:
 
         data = {"query": LIST_RULE_INSTANCES, "flags": {"variableResultSize": True}}
 
-        r = requests.post(
-            url=self.graphql_url, headers=self.headers, json=data, verify=True
+        r = self.session.post(
+            url=self.graphql_url, headers=self.headers, json=data, timeout=60
         ).json()
         results.extend(r["data"]["listRuleInstances"]["questionInstances"])
 
@@ -1220,8 +1228,8 @@ def list_alert_rules(self) -> List[Dict[str, Any]]:
                 "flags": {"variableResultSize": True},
             }
 
-            r = requests.post(
-                url=self.graphql_url, headers=self.headers, json=data, verify=True
+            r = self.session.post(
+                url=self.graphql_url, headers=self.headers, json=data, timeout=60
             ).json()
             results.extend(r["data"]["listRuleInstances"]["questionInstances"])
 
@@ -1233,8 +1241,8 @@ def get_alert_rule_details(self, rule_id: Optional[str] = None) -> Dict[str, Any
 
         data = {"query": LIST_RULE_INSTANCES, "flags": {"variableResultSize": True}}
 
-        r = requests.post(
-            url=self.graphql_url, headers=self.headers, json=data, verify=True
+        r = self.session.post(
+            url=self.graphql_url, headers=self.headers, json=data, timeout=60
         ).json()
         results.extend(r["data"]["listRuleInstances"]["questionInstances"])
 
@@ -1249,8 +1257,8 @@ def get_alert_rule_details(self, rule_id: Optional[str] = None) -> Dict[str, Any
                 "flags": {"variableResultSize": True},
             }
 
-            r = requests.post(
-                url=self.graphql_url, headers=self.headers, json=data, verify=True
+            r = self.session.post(
+                url=self.graphql_url, headers=self.headers, json=data, timeout=60
             ).json()
             results.extend(r["data"]["listRuleInstances"]["questionInstances"])
 
@@ -1597,8 +1605,8 @@ def list_questions(self, search_query: Optional[str] = None, tags: Optional[List
             }
         }
 
-        r = requests.post(
-            url=self.graphql_url, headers=self.headers, json=data, verify=True
+        r = self.session.post(
+            url=self.graphql_url, headers=self.headers, json=data, timeout=60
         ).json()
         results.extend(r["data"]["questions"]["questions"])
 
@@ -1619,8 +1627,8 @@ def list_questions(self, search_query: Optional[str] = None, tags: Optional[List
                 },
             }
 
-            r = requests.post(
-                url=self.graphql_url, headers=self.headers, json=data, verify=True
+            r = self.session.post(
+                url=self.graphql_url, headers=self.headers, json=data, timeout=60
             ).json()
             results.extend(r["data"]["questions"]["questions"])
 
@@ -1925,8 +1933,8 @@ def list_account_parameters(self):
 
         data = {"query": PARAMETER_LIST, "flags": {"variableResultSize": True}}
 
-        r = requests.post(
-            url=self.graphql_url, headers=self.headers, json=data, verify=True
+        r = self.session.post(
+            url=self.graphql_url, headers=self.headers, json=data, timeout=60
         ).json()
         results.extend(r["data"]["parameterList"]["items"])
 
@@ -1940,8 +1948,8 @@ def list_account_parameters(self):
                 "flags": {"variableResultSize": True},
             }
 
-            r = requests.post(
-                url=self.graphql_url, headers=self.headers, json=data, verify=True
+            r = self.session.post(
+                url=self.graphql_url, headers=self.headers, json=data, timeout=60
             ).json()
             results.extend(r["data"]["parameterList"]["items"])
 

setup.py

@@ -5,7 +5,7 @@
 
 setup(
     name="jupiterone",
-    version="2.2.0",
+    version="2.3.0",
     description="A Python client for the JupiterOne API",
     license="MIT License",
     author="JupiterOne",

tests/test_alert_rule_methods.py

@@ -15,7 +15,7 @@ def setup_method(self):
         """Set up test fixtures"""
         self.client = JupiterOneClient(account="test-account", token="test-token")
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_alert_rules(self, mock_post):
         """Test list_alert_rules method"""
         # Mock first page response
@@ -55,7 +55,7 @@ def test_list_alert_rules(self, mock_post):
         assert result[1]["id"] == "rule-2"
         assert mock_post.call_count == 2
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_get_alert_rule_details_found(self, mock_post):
         """Test get_alert_rule_details method - rule found"""
         # Mock response with the target rule
@@ -81,7 +81,7 @@ def test_get_alert_rule_details_found(self, mock_post):
         assert result["id"] == "rule-1"
         assert result["name"] == "Test Rule"
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_get_alert_rule_details_not_found(self, mock_post):
         """Test get_alert_rule_details method - rule not found"""
         # Mock response without the target rule

tests/test_delete_entity.py

@@ -1,5 +1,4 @@
 import json
-import pytest
 import responses
 
 from jupiterone.client import JupiterOneClient

tests/test_list_questions.py

@@ -15,7 +15,7 @@ def setup_method(self):
         """Set up test fixtures"""
         self.client = JupiterOneClient(account="test-account", token="test-token")
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_basic(self, mock_post):
         """Test basic questions listing with single page"""
         # Mock response for single page
@@ -101,7 +101,7 @@ def test_list_questions_basic(self, mock_post):
         assert call_args[1]['json']['query'] == QUESTIONS
         assert call_args[1]['json']['flags']['variableResultSize'] is True
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_with_pagination(self, mock_post):
         """Test questions listing with multiple pages"""
         # Mock first page response
@@ -198,7 +198,7 @@ def test_list_questions_with_pagination(self, mock_post):
         assert second_call[1]['json']['variables']['cursor'] == "cursor-1"
         assert second_call[1]['json']['flags']['variableResultSize'] is True
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_empty_response(self, mock_post):
         """Test questions listing with empty response"""
         # Mock empty response
@@ -227,7 +227,7 @@ def test_list_questions_empty_response(self, mock_post):
         # Verify API call
         mock_post.assert_called_once()
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_with_compliance_data(self, mock_post):
         """Test questions listing with compliance metadata"""
         # Mock response with compliance data
@@ -279,7 +279,7 @@ def test_list_questions_with_compliance_data(self, mock_post):
         assert compliance['requirements'] == ["2.1", "2.2", "2.3"]
         assert compliance['controls'] == ["Data Protection", "Network Security"]
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_with_variables(self, mock_post):
         """Test questions listing with variable definitions"""
         # Mock response with variables
@@ -341,7 +341,7 @@ def test_list_questions_with_variables(self, mock_post):
         assert variables[1]['required'] is False
         assert variables[1]['default'] == "us-east-1"
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_with_polling_intervals(self, mock_post):
         """Test questions listing with different polling intervals"""
         # Mock response with various polling intervals
@@ -413,7 +413,7 @@ def test_list_questions_with_polling_intervals(self, mock_post):
         assert result[1]['showTrend'] is True
         assert result[2]['showTrend'] is False
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_error_handling(self, mock_post):
         """Test questions listing with error handling"""
         # Mock error response
@@ -433,7 +433,7 @@ def test_list_questions_error_handling(self, mock_post):
         with pytest.raises(Exception):
             self.client.list_questions()
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_malformed_response(self, mock_post):
         """Test questions listing with malformed response"""
         # Mock malformed response
@@ -466,7 +466,7 @@ def test_list_questions_malformed_response(self, mock_post):
         # Missing fields should be None or not present
         assert 'title' not in question or question['title'] is None
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_with_search_query(self, mock_post):
         """Test questions listing with search query parameter"""
         # Mock response for search query
@@ -512,7 +512,7 @@ def test_list_questions_with_search_query(self, mock_post):
         call_args = mock_post.call_args
         assert call_args[1]['json']['variables']['searchQuery'] == "security"
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_with_tags_filter(self, mock_post):
         """Test questions listing with tags filter parameter"""
         # Mock response for tags filter
@@ -563,7 +563,7 @@ def test_list_questions_with_tags_filter(self, mock_post):
         call_args = mock_post.call_args
         assert call_args[1]['json']['variables']['tags'] == ["cis", "aws"]
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_with_search_and_tags(self, mock_post):
         """Test questions listing with both search query and tags filter"""
         # Mock response for combined search and tags
@@ -620,7 +620,7 @@ def test_list_questions_with_search_and_tags(self, mock_post):
         assert variables['searchQuery'] == "encryption"
         assert variables['tags'] == ["security", "compliance"]
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_with_pagination_and_filters(self, mock_post):
         """Test questions listing with filters and pagination"""
         # Mock first page response with filters
@@ -709,7 +709,7 @@ def test_list_questions_with_pagination_and_filters(self, mock_post):
         assert second_variables['tags'] == ["security"]
         assert second_variables['cursor'] == "cursor-1"
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_no_parameters(self, mock_post):
         """Test questions listing with no parameters (default behavior)"""
         # Mock response for no parameters
@@ -753,7 +753,7 @@ def test_list_questions_no_parameters(self, mock_post):
         call_args = mock_post.call_args
         assert call_args[1]['json']['variables'] == {}
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_empty_search_results(self, mock_post):
         """Test questions listing with search that returns no results"""
         # Mock empty response for search
@@ -784,7 +784,7 @@ def test_list_questions_empty_search_results(self, mock_post):
         call_args = mock_post.call_args
         assert call_args[1]['json']['variables']['searchQuery'] == "nonexistent"
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions_empty_tags_results(self, mock_post):
         """Test questions listing with tags filter that returns no results"""
         # Mock empty response for tags filter

tests/test_misc_methods.py

@@ -14,7 +14,7 @@ def setup_method(self):
         """Set up test fixtures"""
         self.client = JupiterOneClient(account="test-account", token="test-token")
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_questions(self, mock_post):
         """Test list_questions method"""
         # Mock first page response
@@ -91,7 +91,7 @@ def test_get_parameter_details(self, mock_execute_query):
         assert result == mock_response
         mock_execute_query.assert_called_once()
 
-    @patch('jupiterone.client.requests.post')
+    @patch('requests.Session.post')
     def test_list_account_parameters(self, mock_post):
         """Test list_account_parameters method"""
         # Mock first page response