Skip to content

Commit ebe1ef9

Browse files
SeaBloomsSeaBlooms
authored
Merge pull request #68 from JupiterOne/release/2.3.0
Custom User-Agent header, Migrated all bare requests.post() to self.session.post(), Version bump to 2.3.0
2 parents 20d6a53 + e6e13fe commit ebe1ef9

6 files changed

Lines changed: 45 additions & 38 deletions

File tree

‎jupiterone/client.py‎

Lines changed: 24 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ class JupiterOneClient:
6868

6969
# pylint: disable=too-many-instance-attributes
7070

71+
SDK_VERSION: str = "2.3.0"
7172
DEFAULT_URL: str = "https://graphql.us.jupiterone.io"
7273
SYNC_API_URL: str = "https://api.us.jupiterone.io"
7374

@@ -77,21 +78,28 @@ def __init__(
7778
token: Optional[str] = None,
7879
url: str = DEFAULT_URL,
7980
sync_url: str = SYNC_API_URL,
81+
user_agent: Optional[str] = None,
8082
) -> None:
8183
# Validate inputs
8284
self._validate_constructor_inputs(account, token, url, sync_url)
8385
self.account: Optional[str] = account
8486
self.token: Optional[str] = token
8587
self.graphql_url: str = url
8688
self.sync_url: str = sync_url
89+
90+
base_ua = f"jupiterone-client-python/{self.SDK_VERSION}"
91+
full_ua = f"{base_ua} {user_agent}" if user_agent else base_ua
92+
8793
self.headers: Dict[str, str] = {
8894
"Authorization": "Bearer {}".format(self.token or ""),
8995
"JupiterOne-Account": self.account or "",
9096
"Content-Type": "application/json",
97+
"User-Agent": full_ua,
9198
}
9299

93100
# Initialize session with retry logic
94101
self.session: requests.Session = requests.Session()
102+
self.session.headers["User-Agent"] = full_ua
95103
retries = Retry(
96104
total=5,
97105
backoff_factor=1,
@@ -1204,8 +1212,8 @@ def list_alert_rules(self) -> List[Dict[str, Any]]:
12041212

12051213
data = {"query": LIST_RULE_INSTANCES, "flags": {"variableResultSize": True}}
12061214

1207-
r = requests.post(
1208-
url=self.graphql_url, headers=self.headers, json=data, verify=True
1215+
r = self.session.post(
1216+
url=self.graphql_url, headers=self.headers, json=data, timeout=60
12091217
).json()
12101218
results.extend(r["data"]["listRuleInstances"]["questionInstances"])
12111219

@@ -1220,8 +1228,8 @@ def list_alert_rules(self) -> List[Dict[str, Any]]:
12201228
"flags": {"variableResultSize": True},
12211229
}
12221230

1223-
r = requests.post(
1224-
url=self.graphql_url, headers=self.headers, json=data, verify=True
1231+
r = self.session.post(
1232+
url=self.graphql_url, headers=self.headers, json=data, timeout=60
12251233
).json()
12261234
results.extend(r["data"]["listRuleInstances"]["questionInstances"])
12271235

@@ -1233,8 +1241,8 @@ def get_alert_rule_details(self, rule_id: Optional[str] = None) -> Dict[str, Any
12331241

12341242
data = {"query": LIST_RULE_INSTANCES, "flags": {"variableResultSize": True}}
12351243

1236-
r = requests.post(
1237-
url=self.graphql_url, headers=self.headers, json=data, verify=True
1244+
r = self.session.post(
1245+
url=self.graphql_url, headers=self.headers, json=data, timeout=60
12381246
).json()
12391247
results.extend(r["data"]["listRuleInstances"]["questionInstances"])
12401248

@@ -1249,8 +1257,8 @@ def get_alert_rule_details(self, rule_id: Optional[str] = None) -> Dict[str, Any
12491257
"flags": {"variableResultSize": True},
12501258
}
12511259

1252-
r = requests.post(
1253-
url=self.graphql_url, headers=self.headers, json=data, verify=True
1260+
r = self.session.post(
1261+
url=self.graphql_url, headers=self.headers, json=data, timeout=60
12541262
).json()
12551263
results.extend(r["data"]["listRuleInstances"]["questionInstances"])
12561264

@@ -1597,8 +1605,8 @@ def list_questions(self, search_query: Optional[str] = None, tags: Optional[List
15971605
}
15981606
}
15991607

1600-
r = requests.post(
1601-
url=self.graphql_url, headers=self.headers, json=data, verify=True
1608+
r = self.session.post(
1609+
url=self.graphql_url, headers=self.headers, json=data, timeout=60
16021610
).json()
16031611
results.extend(r["data"]["questions"]["questions"])
16041612

@@ -1619,8 +1627,8 @@ def list_questions(self, search_query: Optional[str] = None, tags: Optional[List
16191627
},
16201628
}
16211629

1622-
r = requests.post(
1623-
url=self.graphql_url, headers=self.headers, json=data, verify=True
1630+
r = self.session.post(
1631+
url=self.graphql_url, headers=self.headers, json=data, timeout=60
16241632
).json()
16251633
results.extend(r["data"]["questions"]["questions"])
16261634

@@ -1925,8 +1933,8 @@ def list_account_parameters(self):
19251933

19261934
data = {"query": PARAMETER_LIST, "flags": {"variableResultSize": True}}
19271935

1928-
r = requests.post(
1929-
url=self.graphql_url, headers=self.headers, json=data, verify=True
1936+
r = self.session.post(
1937+
url=self.graphql_url, headers=self.headers, json=data, timeout=60
19301938
).json()
19311939
results.extend(r["data"]["parameterList"]["items"])
19321940

@@ -1940,8 +1948,8 @@ def list_account_parameters(self):
19401948
"flags": {"variableResultSize": True},
19411949
}
19421950

1943-
r = requests.post(
1944-
url=self.graphql_url, headers=self.headers, json=data, verify=True
1951+
r = self.session.post(
1952+
url=self.graphql_url, headers=self.headers, json=data, timeout=60
19451953
).json()
19461954
results.extend(r["data"]["parameterList"]["items"])
19471955

‎setup.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
setup(
77
name="jupiterone",
8-
version="2.2.0",
8+
version="2.3.0",
99
description="A Python client for the JupiterOne API",
1010
license="MIT License",
1111
author="JupiterOne",

‎tests/test_alert_rule_methods.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ def setup_method(self):
1515
"""Set up test fixtures"""
1616
self.client = JupiterOneClient(account="test-account", token="test-token")
1717

18-
@patch('jupiterone.client.requests.post')
18+
@patch('requests.Session.post')
1919
def test_list_alert_rules(self, mock_post):
2020
"""Test list_alert_rules method"""
2121
# Mock first page response
@@ -55,7 +55,7 @@ def test_list_alert_rules(self, mock_post):
5555
assert result[1]["id"] == "rule-2"
5656
assert mock_post.call_count == 2
5757

58-
@patch('jupiterone.client.requests.post')
58+
@patch('requests.Session.post')
5959
def test_get_alert_rule_details_found(self, mock_post):
6060
"""Test get_alert_rule_details method - rule found"""
6161
# Mock response with the target rule
@@ -81,7 +81,7 @@ def test_get_alert_rule_details_found(self, mock_post):
8181
assert result["id"] == "rule-1"
8282
assert result["name"] == "Test Rule"
8383

84-
@patch('jupiterone.client.requests.post')
84+
@patch('requests.Session.post')
8585
def test_get_alert_rule_details_not_found(self, mock_post):
8686
"""Test get_alert_rule_details method - rule not found"""
8787
# Mock response without the target rule

‎tests/test_delete_entity.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
import json
2-
import pytest
32
import responses
43

54
from jupiterone.client import JupiterOneClient

‎tests/test_list_questions.py‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ def setup_method(self):
1515
"""Set up test fixtures"""
1616
self.client = JupiterOneClient(account="test-account", token="test-token")
1717

18-
@patch('jupiterone.client.requests.post')
18+
@patch('requests.Session.post')
1919
def test_list_questions_basic(self, mock_post):
2020
"""Test basic questions listing with single page"""
2121
# Mock response for single page
@@ -101,7 +101,7 @@ def test_list_questions_basic(self, mock_post):
101101
assert call_args[1]['json']['query'] == QUESTIONS
102102
assert call_args[1]['json']['flags']['variableResultSize'] is True
103103

104-
@patch('jupiterone.client.requests.post')
104+
@patch('requests.Session.post')
105105
def test_list_questions_with_pagination(self, mock_post):
106106
"""Test questions listing with multiple pages"""
107107
# Mock first page response
@@ -198,7 +198,7 @@ def test_list_questions_with_pagination(self, mock_post):
198198
assert second_call[1]['json']['variables']['cursor'] == "cursor-1"
199199
assert second_call[1]['json']['flags']['variableResultSize'] is True
200200

201-
@patch('jupiterone.client.requests.post')
201+
@patch('requests.Session.post')
202202
def test_list_questions_empty_response(self, mock_post):
203203
"""Test questions listing with empty response"""
204204
# Mock empty response
@@ -227,7 +227,7 @@ def test_list_questions_empty_response(self, mock_post):
227227
# Verify API call
228228
mock_post.assert_called_once()
229229

230-
@patch('jupiterone.client.requests.post')
230+
@patch('requests.Session.post')
231231
def test_list_questions_with_compliance_data(self, mock_post):
232232
"""Test questions listing with compliance metadata"""
233233
# Mock response with compliance data
@@ -279,7 +279,7 @@ def test_list_questions_with_compliance_data(self, mock_post):
279279
assert compliance['requirements'] == ["2.1", "2.2", "2.3"]
280280
assert compliance['controls'] == ["Data Protection", "Network Security"]
281281

282-
@patch('jupiterone.client.requests.post')
282+
@patch('requests.Session.post')
283283
def test_list_questions_with_variables(self, mock_post):
284284
"""Test questions listing with variable definitions"""
285285
# Mock response with variables
@@ -341,7 +341,7 @@ def test_list_questions_with_variables(self, mock_post):
341341
assert variables[1]['required'] is False
342342
assert variables[1]['default'] == "us-east-1"
343343

344-
@patch('jupiterone.client.requests.post')
344+
@patch('requests.Session.post')
345345
def test_list_questions_with_polling_intervals(self, mock_post):
346346
"""Test questions listing with different polling intervals"""
347347
# Mock response with various polling intervals
@@ -413,7 +413,7 @@ def test_list_questions_with_polling_intervals(self, mock_post):
413413
assert result[1]['showTrend'] is True
414414
assert result[2]['showTrend'] is False
415415

416-
@patch('jupiterone.client.requests.post')
416+
@patch('requests.Session.post')
417417
def test_list_questions_error_handling(self, mock_post):
418418
"""Test questions listing with error handling"""
419419
# Mock error response
@@ -433,7 +433,7 @@ def test_list_questions_error_handling(self, mock_post):
433433
with pytest.raises(Exception):
434434
self.client.list_questions()
435435

436-
@patch('jupiterone.client.requests.post')
436+
@patch('requests.Session.post')
437437
def test_list_questions_malformed_response(self, mock_post):
438438
"""Test questions listing with malformed response"""
439439
# Mock malformed response
@@ -466,7 +466,7 @@ def test_list_questions_malformed_response(self, mock_post):
466466
# Missing fields should be None or not present
467467
assert 'title' not in question or question['title'] is None
468468

469-
@patch('jupiterone.client.requests.post')
469+
@patch('requests.Session.post')
470470
def test_list_questions_with_search_query(self, mock_post):
471471
"""Test questions listing with search query parameter"""
472472
# Mock response for search query
@@ -512,7 +512,7 @@ def test_list_questions_with_search_query(self, mock_post):
512512
call_args = mock_post.call_args
513513
assert call_args[1]['json']['variables']['searchQuery'] == "security"
514514

515-
@patch('jupiterone.client.requests.post')
515+
@patch('requests.Session.post')
516516
def test_list_questions_with_tags_filter(self, mock_post):
517517
"""Test questions listing with tags filter parameter"""
518518
# Mock response for tags filter
@@ -563,7 +563,7 @@ def test_list_questions_with_tags_filter(self, mock_post):
563563
call_args = mock_post.call_args
564564
assert call_args[1]['json']['variables']['tags'] == ["cis", "aws"]
565565

566-
@patch('jupiterone.client.requests.post')
566+
@patch('requests.Session.post')
567567
def test_list_questions_with_search_and_tags(self, mock_post):
568568
"""Test questions listing with both search query and tags filter"""
569569
# Mock response for combined search and tags
@@ -620,7 +620,7 @@ def test_list_questions_with_search_and_tags(self, mock_post):
620620
assert variables['searchQuery'] == "encryption"
621621
assert variables['tags'] == ["security", "compliance"]
622622

623-
@patch('jupiterone.client.requests.post')
623+
@patch('requests.Session.post')
624624
def test_list_questions_with_pagination_and_filters(self, mock_post):
625625
"""Test questions listing with filters and pagination"""
626626
# Mock first page response with filters
@@ -709,7 +709,7 @@ def test_list_questions_with_pagination_and_filters(self, mock_post):
709709
assert second_variables['tags'] == ["security"]
710710
assert second_variables['cursor'] == "cursor-1"
711711

712-
@patch('jupiterone.client.requests.post')
712+
@patch('requests.Session.post')
713713
def test_list_questions_no_parameters(self, mock_post):
714714
"""Test questions listing with no parameters (default behavior)"""
715715
# Mock response for no parameters
@@ -753,7 +753,7 @@ def test_list_questions_no_parameters(self, mock_post):
753753
call_args = mock_post.call_args
754754
assert call_args[1]['json']['variables'] == {}
755755

756-
@patch('jupiterone.client.requests.post')
756+
@patch('requests.Session.post')
757757
def test_list_questions_empty_search_results(self, mock_post):
758758
"""Test questions listing with search that returns no results"""
759759
# Mock empty response for search
@@ -784,7 +784,7 @@ def test_list_questions_empty_search_results(self, mock_post):
784784
call_args = mock_post.call_args
785785
assert call_args[1]['json']['variables']['searchQuery'] == "nonexistent"
786786

787-
@patch('jupiterone.client.requests.post')
787+
@patch('requests.Session.post')
788788
def test_list_questions_empty_tags_results(self, mock_post):
789789
"""Test questions listing with tags filter that returns no results"""
790790
# Mock empty response for tags filter

‎tests/test_misc_methods.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ def setup_method(self):
1414
"""Set up test fixtures"""
1515
self.client = JupiterOneClient(account="test-account", token="test-token")
1616

17-
@patch('jupiterone.client.requests.post')
17+
@patch('requests.Session.post')
1818
def test_list_questions(self, mock_post):
1919
"""Test list_questions method"""
2020
# Mock first page response
@@ -91,7 +91,7 @@ def test_get_parameter_details(self, mock_execute_query):
9191
assert result == mock_response
9292
mock_execute_query.assert_called_once()
9393

94-
@patch('jupiterone.client.requests.post')
94+
@patch('requests.Session.post')
9595
def test_list_account_parameters(self, mock_post):
9696
"""Test list_account_parameters method"""
9797
# Mock first page response

0 commit comments

Comments
 (0)