regscale-cli 6.26.0.0__py3-none-any.whl → 6.27.0.1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of regscale-cli might be problematic. Click here for more details.
- regscale/_version.py +1 -1
- regscale/core/app/application.py +1 -1
- regscale/core/app/internal/evidence.py +419 -2
- regscale/dev/code_gen.py +24 -20
- regscale/integrations/commercial/__init__.py +0 -1
- regscale/integrations/commercial/jira.py +367 -126
- regscale/integrations/commercial/qualys/__init__.py +7 -8
- regscale/integrations/commercial/qualys/scanner.py +8 -3
- regscale/integrations/commercial/synqly/assets.py +17 -0
- regscale/integrations/commercial/synqly/vulnerabilities.py +45 -28
- regscale/integrations/commercial/tenablev2/cis_parsers.py +453 -0
- regscale/integrations/commercial/tenablev2/cis_scanner.py +447 -0
- regscale/integrations/commercial/tenablev2/commands.py +142 -1
- regscale/integrations/commercial/tenablev2/scanner.py +0 -1
- regscale/integrations/commercial/tenablev2/stig_parsers.py +113 -57
- regscale/integrations/commercial/wizv2/WizDataMixin.py +1 -1
- regscale/integrations/commercial/wizv2/click.py +44 -59
- regscale/integrations/commercial/wizv2/compliance/__init__.py +15 -0
- regscale/integrations/commercial/wizv2/{policy_compliance_helpers.py → compliance/helpers.py} +78 -60
- regscale/integrations/commercial/wizv2/compliance_report.py +10 -9
- regscale/integrations/commercial/wizv2/core/__init__.py +133 -0
- regscale/integrations/commercial/wizv2/{async_client.py → core/client.py} +3 -3
- regscale/integrations/commercial/wizv2/{constants.py → core/constants.py} +1 -17
- regscale/integrations/commercial/wizv2/core/file_operations.py +237 -0
- regscale/integrations/commercial/wizv2/fetchers/__init__.py +11 -0
- regscale/integrations/commercial/wizv2/{data_fetcher.py → fetchers/policy_assessment.py} +5 -9
- regscale/integrations/commercial/wizv2/issue.py +1 -1
- regscale/integrations/commercial/wizv2/models/__init__.py +0 -0
- regscale/integrations/commercial/wizv2/parsers/__init__.py +34 -0
- regscale/integrations/commercial/wizv2/{parsers.py → parsers/main.py} +1 -1
- regscale/integrations/commercial/wizv2/processors/__init__.py +11 -0
- regscale/integrations/commercial/wizv2/{finding_processor.py → processors/finding.py} +1 -1
- regscale/integrations/commercial/wizv2/reports.py +1 -1
- regscale/integrations/commercial/wizv2/sbom.py +1 -1
- regscale/integrations/commercial/wizv2/scanner.py +40 -100
- regscale/integrations/commercial/wizv2/utils/__init__.py +48 -0
- regscale/integrations/commercial/wizv2/{utils.py → utils/main.py} +116 -61
- regscale/integrations/commercial/wizv2/variables.py +89 -3
- regscale/integrations/compliance_integration.py +0 -46
- regscale/integrations/control_matcher.py +22 -3
- regscale/integrations/due_date_handler.py +14 -8
- regscale/integrations/public/fedramp/docx_parser.py +10 -1
- regscale/integrations/public/fedramp/fedramp_cis_crm.py +393 -340
- regscale/integrations/public/fedramp/fedramp_five.py +1 -1
- regscale/integrations/scanner_integration.py +127 -57
- regscale/models/integration_models/cisa_kev_data.json +132 -9
- regscale/models/integration_models/qualys.py +3 -4
- regscale/models/integration_models/synqly_models/capabilities.json +1 -1
- regscale/models/integration_models/synqly_models/connectors/vulnerabilities.py +24 -7
- regscale/models/integration_models/synqly_models/synqly_model.py +8 -1
- regscale/models/regscale_models/control_implementation.py +1 -1
- regscale/models/regscale_models/issue.py +0 -1
- {regscale_cli-6.26.0.0.dist-info → regscale_cli-6.27.0.1.dist-info}/METADATA +1 -17
- {regscale_cli-6.26.0.0.dist-info → regscale_cli-6.27.0.1.dist-info}/RECORD +94 -61
- tests/regscale/integrations/commercial/test_jira.py +481 -91
- tests/regscale/integrations/commercial/test_wiz.py +96 -200
- tests/regscale/integrations/commercial/wizv2/__init__.py +1 -1
- tests/regscale/integrations/commercial/wizv2/compliance/__init__.py +1 -0
- tests/regscale/integrations/commercial/wizv2/compliance/test_helpers.py +903 -0
- tests/regscale/integrations/commercial/wizv2/core/__init__.py +1 -0
- tests/regscale/integrations/commercial/wizv2/core/test_auth.py +701 -0
- tests/regscale/integrations/commercial/wizv2/core/test_client.py +1037 -0
- tests/regscale/integrations/commercial/wizv2/core/test_file_operations.py +989 -0
- tests/regscale/integrations/commercial/wizv2/fetchers/__init__.py +1 -0
- tests/regscale/integrations/commercial/wizv2/fetchers/test_policy_assessment.py +805 -0
- tests/regscale/integrations/commercial/wizv2/parsers/__init__.py +1 -0
- tests/regscale/integrations/commercial/wizv2/parsers/test_main.py +1153 -0
- tests/regscale/integrations/commercial/wizv2/processors/__init__.py +1 -0
- tests/regscale/integrations/commercial/wizv2/processors/test_finding.py +671 -0
- tests/regscale/integrations/commercial/wizv2/test_WizDataMixin.py +537 -0
- tests/regscale/integrations/commercial/wizv2/test_click_comprehensive.py +851 -0
- tests/regscale/integrations/commercial/wizv2/test_compliance_report_comprehensive.py +910 -0
- tests/regscale/integrations/commercial/wizv2/test_file_cleanup.py +283 -0
- tests/regscale/integrations/commercial/wizv2/test_file_operations.py +260 -0
- tests/regscale/integrations/commercial/wizv2/test_issue.py +1 -1
- tests/regscale/integrations/commercial/wizv2/test_issue_comprehensive.py +1203 -0
- tests/regscale/integrations/commercial/wizv2/test_reports.py +497 -0
- tests/regscale/integrations/commercial/wizv2/test_sbom.py +643 -0
- tests/regscale/integrations/commercial/wizv2/test_scanner_comprehensive.py +805 -0
- tests/regscale/integrations/commercial/wizv2/test_wiz_click_client_id.py +1 -1
- tests/regscale/integrations/commercial/wizv2/test_wiz_compliance_report.py +72 -29
- tests/regscale/integrations/commercial/wizv2/test_wiz_findings_comprehensive.py +364 -0
- tests/regscale/integrations/commercial/wizv2/test_wiz_inventory_comprehensive.py +644 -0
- tests/regscale/integrations/commercial/wizv2/test_wizv2.py +946 -78
- tests/regscale/integrations/commercial/wizv2/test_wizv2_utils.py +97 -202
- tests/regscale/integrations/commercial/wizv2/utils/__init__.py +1 -0
- tests/regscale/integrations/commercial/wizv2/utils/test_main.py +1523 -0
- tests/regscale/integrations/public/test_fedramp.py +301 -0
- tests/regscale/integrations/test_control_matcher.py +83 -0
- regscale/integrations/commercial/wizv2/policy_compliance.py +0 -3543
- tests/regscale/integrations/commercial/wizv2/test_wiz_policy_compliance.py +0 -750
- /regscale/integrations/commercial/wizv2/{wiz_auth.py → core/auth.py} +0 -0
- {regscale_cli-6.26.0.0.dist-info → regscale_cli-6.27.0.1.dist-info}/LICENSE +0 -0
- {regscale_cli-6.26.0.0.dist-info → regscale_cli-6.27.0.1.dist-info}/WHEEL +0 -0
- {regscale_cli-6.26.0.0.dist-info → regscale_cli-6.27.0.1.dist-info}/entry_points.txt +0 -0
- {regscale_cli-6.26.0.0.dist-info → regscale_cli-6.27.0.1.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,1037 @@
|
|
|
1
|
+
#!/usr/bin/env python3
|
|
2
|
+
# -*- coding: utf-8 -*-
|
|
3
|
+
"""Comprehensive unit tests for Wiz AsyncGraphQLClient Module"""
|
|
4
|
+
|
|
5
|
+
import asyncio
|
|
6
|
+
import logging
|
|
7
|
+
from unittest.mock import AsyncMock, MagicMock, Mock, call, patch
|
|
8
|
+
|
|
9
|
+
import anyio
|
|
10
|
+
import httpx
|
|
11
|
+
import pytest
|
|
12
|
+
|
|
13
|
+
from regscale.integrations.commercial.wizv2.core.client import (
|
|
14
|
+
AsyncWizGraphQLClient,
|
|
15
|
+
run_async_queries,
|
|
16
|
+
)
|
|
17
|
+
from regscale.integrations.commercial.wizv2.core.constants import WizVulnerabilityType
|
|
18
|
+
|
|
19
|
+
PATH = "regscale.integrations.commercial.wizv2.core.client"
|
|
20
|
+
|
|
21
|
+
logger = logging.getLogger("regscale")
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class TestAsyncWizGraphQLClientInit:
|
|
25
|
+
"""Test AsyncWizGraphQLClient initialization"""
|
|
26
|
+
|
|
27
|
+
def test_init_with_defaults(self):
|
|
28
|
+
"""Test initialization with default parameters"""
|
|
29
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
30
|
+
|
|
31
|
+
assert client.endpoint == "https://api.wiz.io/graphql"
|
|
32
|
+
assert client.headers == {}
|
|
33
|
+
assert client.timeout == 30.0
|
|
34
|
+
assert client.max_concurrent == 5
|
|
35
|
+
assert client._semaphore is not None
|
|
36
|
+
|
|
37
|
+
def test_init_with_custom_parameters(self):
|
|
38
|
+
"""Test initialization with custom parameters"""
|
|
39
|
+
headers = {"Authorization": "Bearer test-token"}
|
|
40
|
+
client = AsyncWizGraphQLClient(
|
|
41
|
+
endpoint="https://custom.wiz.io/graphql",
|
|
42
|
+
headers=headers,
|
|
43
|
+
timeout=60.0,
|
|
44
|
+
max_concurrent=10,
|
|
45
|
+
)
|
|
46
|
+
|
|
47
|
+
assert client.endpoint == "https://custom.wiz.io/graphql"
|
|
48
|
+
assert client.headers == headers
|
|
49
|
+
assert client.timeout == 60.0
|
|
50
|
+
assert client.max_concurrent == 10
|
|
51
|
+
|
|
52
|
+
def test_init_with_none_headers(self):
|
|
53
|
+
"""Test initialization with None headers defaults to empty dict"""
|
|
54
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql", headers=None)
|
|
55
|
+
|
|
56
|
+
assert client.headers == {}
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
class TestAsyncWizGraphQLClientExecuteQuery:
|
|
60
|
+
"""Test AsyncWizGraphQLClient.execute_query method"""
|
|
61
|
+
|
|
62
|
+
@pytest.mark.asyncio
|
|
63
|
+
async def test_execute_query_success(self):
|
|
64
|
+
"""Test successful query execution"""
|
|
65
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
66
|
+
|
|
67
|
+
mock_response = Mock()
|
|
68
|
+
mock_response.is_success = True
|
|
69
|
+
mock_response.json.return_value = {"data": {"result": "success"}}
|
|
70
|
+
|
|
71
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
72
|
+
mock_client = AsyncMock()
|
|
73
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
74
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
75
|
+
|
|
76
|
+
result = await client.execute_query(query="query { test }", variables={"key": "value"})
|
|
77
|
+
|
|
78
|
+
assert result == {"result": "success"}
|
|
79
|
+
mock_client.post.assert_called_once()
|
|
80
|
+
|
|
81
|
+
@pytest.mark.asyncio
|
|
82
|
+
async def test_execute_query_with_progress_callback(self):
|
|
83
|
+
"""Test query execution with progress callback"""
|
|
84
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
85
|
+
|
|
86
|
+
mock_response = Mock()
|
|
87
|
+
mock_response.is_success = True
|
|
88
|
+
mock_response.json.return_value = {"data": {"result": "success"}}
|
|
89
|
+
|
|
90
|
+
progress_calls = []
|
|
91
|
+
|
|
92
|
+
def progress_callback(task_name, status, extra_data=None):
|
|
93
|
+
progress_calls.append((task_name, status))
|
|
94
|
+
|
|
95
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
96
|
+
mock_client = AsyncMock()
|
|
97
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
98
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
99
|
+
|
|
100
|
+
result = await client.execute_query(
|
|
101
|
+
query="query { test }",
|
|
102
|
+
variables={"key": "value"},
|
|
103
|
+
progress_callback=progress_callback,
|
|
104
|
+
task_name="Test Query",
|
|
105
|
+
)
|
|
106
|
+
|
|
107
|
+
assert result == {"result": "success"}
|
|
108
|
+
assert ("Test Query", "starting") in progress_calls
|
|
109
|
+
assert ("Test Query", "requesting") in progress_calls
|
|
110
|
+
assert ("Test Query", "processing") in progress_calls
|
|
111
|
+
assert ("Test Query", "completed") in progress_calls
|
|
112
|
+
|
|
113
|
+
@pytest.mark.asyncio
|
|
114
|
+
async def test_execute_query_non_success_response(self):
|
|
115
|
+
"""Test query execution with non-success HTTP response"""
|
|
116
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
117
|
+
|
|
118
|
+
mock_response = Mock()
|
|
119
|
+
mock_response.is_success = False
|
|
120
|
+
mock_response.status_code = 500
|
|
121
|
+
mock_response.text = "Internal Server Error"
|
|
122
|
+
|
|
123
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
124
|
+
mock_client = AsyncMock()
|
|
125
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
126
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
127
|
+
|
|
128
|
+
with patch(f"{PATH}.error_and_exit") as mock_error_exit:
|
|
129
|
+
mock_error_exit.side_effect = SystemExit(1)
|
|
130
|
+
|
|
131
|
+
with pytest.raises(SystemExit):
|
|
132
|
+
await client.execute_query(query="query { test }")
|
|
133
|
+
|
|
134
|
+
mock_error_exit.assert_called_once()
|
|
135
|
+
error_message = mock_error_exit.call_args[0][0]
|
|
136
|
+
assert "500" in error_message
|
|
137
|
+
|
|
138
|
+
@pytest.mark.asyncio
|
|
139
|
+
async def test_execute_query_with_graphql_errors(self):
|
|
140
|
+
"""Test query execution with GraphQL errors in response"""
|
|
141
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
142
|
+
|
|
143
|
+
mock_response = Mock()
|
|
144
|
+
mock_response.is_success = True
|
|
145
|
+
mock_response.json.return_value = {"errors": [{"message": "GraphQL error occurred"}]}
|
|
146
|
+
|
|
147
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
148
|
+
mock_client = AsyncMock()
|
|
149
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
150
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
151
|
+
|
|
152
|
+
with patch(f"{PATH}.error_and_exit") as mock_error_exit:
|
|
153
|
+
mock_error_exit.side_effect = SystemExit(1)
|
|
154
|
+
|
|
155
|
+
with pytest.raises(SystemExit):
|
|
156
|
+
await client.execute_query(query="query { test }")
|
|
157
|
+
|
|
158
|
+
mock_error_exit.assert_called_once()
|
|
159
|
+
error_message = mock_error_exit.call_args[0][0]
|
|
160
|
+
assert "GraphQL errors" in error_message
|
|
161
|
+
|
|
162
|
+
@pytest.mark.asyncio
|
|
163
|
+
async def test_execute_query_http_error(self):
|
|
164
|
+
"""Test query execution with HTTP error"""
|
|
165
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
166
|
+
|
|
167
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
168
|
+
mock_client = AsyncMock()
|
|
169
|
+
mock_client.post = AsyncMock(side_effect=httpx.HTTPError("Connection error"))
|
|
170
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
171
|
+
|
|
172
|
+
with patch(f"{PATH}.error_and_exit") as mock_error_exit:
|
|
173
|
+
mock_error_exit.side_effect = SystemExit(1)
|
|
174
|
+
|
|
175
|
+
with pytest.raises(SystemExit):
|
|
176
|
+
await client.execute_query(query="query { test }")
|
|
177
|
+
|
|
178
|
+
mock_error_exit.assert_called_once()
|
|
179
|
+
error_message = mock_error_exit.call_args[0][0]
|
|
180
|
+
assert "HTTP error" in error_message
|
|
181
|
+
|
|
182
|
+
@pytest.mark.asyncio
|
|
183
|
+
async def test_execute_query_general_exception(self):
|
|
184
|
+
"""Test query execution with general exception"""
|
|
185
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
186
|
+
|
|
187
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
188
|
+
mock_client = AsyncMock()
|
|
189
|
+
mock_client.post = AsyncMock(side_effect=Exception("Unexpected error"))
|
|
190
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
191
|
+
|
|
192
|
+
with patch(f"{PATH}.error_and_exit") as mock_error_exit:
|
|
193
|
+
mock_error_exit.side_effect = SystemExit(1)
|
|
194
|
+
|
|
195
|
+
with pytest.raises(SystemExit):
|
|
196
|
+
await client.execute_query(query="query { test }")
|
|
197
|
+
|
|
198
|
+
mock_error_exit.assert_called_once()
|
|
199
|
+
error_message = mock_error_exit.call_args[0][0]
|
|
200
|
+
assert "Error in" in error_message
|
|
201
|
+
|
|
202
|
+
@pytest.mark.asyncio
|
|
203
|
+
async def test_execute_query_with_ssl_verify_false(self):
|
|
204
|
+
"""Test query execution with SSL verification disabled"""
|
|
205
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
206
|
+
|
|
207
|
+
mock_response = Mock()
|
|
208
|
+
mock_response.is_success = True
|
|
209
|
+
mock_response.json.return_value = {"data": {"result": "success"}}
|
|
210
|
+
|
|
211
|
+
with patch(f"{PATH}.ScannerVariables") as mock_scanner_vars:
|
|
212
|
+
mock_scanner_vars.sslVerify = False
|
|
213
|
+
|
|
214
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
215
|
+
mock_client = AsyncMock()
|
|
216
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
217
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
218
|
+
|
|
219
|
+
result = await client.execute_query(query="query { test }")
|
|
220
|
+
|
|
221
|
+
assert result == {"result": "success"}
|
|
222
|
+
mock_client_class.assert_called_once()
|
|
223
|
+
call_kwargs = mock_client_class.call_args[1]
|
|
224
|
+
assert call_kwargs["verify"] is False
|
|
225
|
+
|
|
226
|
+
@pytest.mark.asyncio
|
|
227
|
+
async def test_execute_query_empty_response_data(self):
|
|
228
|
+
"""Test query execution with empty data in response"""
|
|
229
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
230
|
+
|
|
231
|
+
mock_response = Mock()
|
|
232
|
+
mock_response.is_success = True
|
|
233
|
+
mock_response.json.return_value = {}
|
|
234
|
+
|
|
235
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
236
|
+
mock_client = AsyncMock()
|
|
237
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
238
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
239
|
+
|
|
240
|
+
result = await client.execute_query(query="query { test }")
|
|
241
|
+
|
|
242
|
+
assert result == {}
|
|
243
|
+
|
|
244
|
+
@pytest.mark.asyncio
|
|
245
|
+
async def test_execute_query_with_empty_variables(self):
|
|
246
|
+
"""Test query execution with empty variables"""
|
|
247
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
248
|
+
|
|
249
|
+
mock_response = Mock()
|
|
250
|
+
mock_response.is_success = True
|
|
251
|
+
mock_response.json.return_value = {"data": {"result": "success"}}
|
|
252
|
+
|
|
253
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
254
|
+
mock_client = AsyncMock()
|
|
255
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
256
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
257
|
+
|
|
258
|
+
result = await client.execute_query(query="query { test }", variables=None)
|
|
259
|
+
|
|
260
|
+
assert result == {"result": "success"}
|
|
261
|
+
call_args = mock_client.post.call_args
|
|
262
|
+
payload = call_args[1]["json"]
|
|
263
|
+
assert payload["variables"] == {}
|
|
264
|
+
|
|
265
|
+
|
|
266
|
+
class TestAsyncWizGraphQLClientExecutePaginatedQuery:
|
|
267
|
+
"""Test AsyncWizGraphQLClient.execute_paginated_query method"""
|
|
268
|
+
|
|
269
|
+
@pytest.mark.asyncio
|
|
270
|
+
async def test_execute_paginated_query_single_page(self):
|
|
271
|
+
"""Test paginated query with single page"""
|
|
272
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
273
|
+
|
|
274
|
+
mock_response = Mock()
|
|
275
|
+
mock_response.is_success = True
|
|
276
|
+
mock_response.json.return_value = {
|
|
277
|
+
"data": {
|
|
278
|
+
"vulnerabilities": {
|
|
279
|
+
"nodes": [{"id": "1"}, {"id": "2"}],
|
|
280
|
+
"pageInfo": {"hasNextPage": False, "endCursor": None},
|
|
281
|
+
}
|
|
282
|
+
}
|
|
283
|
+
}
|
|
284
|
+
|
|
285
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
286
|
+
mock_client = AsyncMock()
|
|
287
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
288
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
289
|
+
|
|
290
|
+
result = await client.execute_paginated_query(
|
|
291
|
+
query="query { vulnerabilities }",
|
|
292
|
+
variables={"first": 100},
|
|
293
|
+
topic_key="vulnerabilities",
|
|
294
|
+
)
|
|
295
|
+
|
|
296
|
+
assert len(result) == 2
|
|
297
|
+
assert result[0]["id"] == "1"
|
|
298
|
+
assert result[1]["id"] == "2"
|
|
299
|
+
|
|
300
|
+
@pytest.mark.asyncio
|
|
301
|
+
async def test_execute_paginated_query_multiple_pages(self):
|
|
302
|
+
"""Test paginated query with multiple pages"""
|
|
303
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
304
|
+
|
|
305
|
+
# First page response
|
|
306
|
+
mock_response_1 = Mock()
|
|
307
|
+
mock_response_1.is_success = True
|
|
308
|
+
mock_response_1.json.return_value = {
|
|
309
|
+
"data": {
|
|
310
|
+
"vulnerabilities": {
|
|
311
|
+
"nodes": [{"id": "1"}, {"id": "2"}],
|
|
312
|
+
"pageInfo": {"hasNextPage": True, "endCursor": "cursor1"},
|
|
313
|
+
}
|
|
314
|
+
}
|
|
315
|
+
}
|
|
316
|
+
|
|
317
|
+
# Second page response
|
|
318
|
+
mock_response_2 = Mock()
|
|
319
|
+
mock_response_2.is_success = True
|
|
320
|
+
mock_response_2.json.return_value = {
|
|
321
|
+
"data": {
|
|
322
|
+
"vulnerabilities": {
|
|
323
|
+
"nodes": [{"id": "3"}, {"id": "4"}],
|
|
324
|
+
"pageInfo": {"hasNextPage": False, "endCursor": None},
|
|
325
|
+
}
|
|
326
|
+
}
|
|
327
|
+
}
|
|
328
|
+
|
|
329
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
330
|
+
mock_client = AsyncMock()
|
|
331
|
+
mock_client.post = AsyncMock(side_effect=[mock_response_1, mock_response_2])
|
|
332
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
333
|
+
|
|
334
|
+
result = await client.execute_paginated_query(
|
|
335
|
+
query="query { vulnerabilities }",
|
|
336
|
+
variables={"first": 100},
|
|
337
|
+
topic_key="vulnerabilities",
|
|
338
|
+
)
|
|
339
|
+
|
|
340
|
+
assert len(result) == 4
|
|
341
|
+
assert result[0]["id"] == "1"
|
|
342
|
+
assert result[3]["id"] == "4"
|
|
343
|
+
assert mock_client.post.call_count == 2
|
|
344
|
+
|
|
345
|
+
@pytest.mark.asyncio
|
|
346
|
+
async def test_execute_paginated_query_with_progress_callback(self):
|
|
347
|
+
"""Test paginated query with progress callback"""
|
|
348
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
349
|
+
|
|
350
|
+
mock_response = Mock()
|
|
351
|
+
mock_response.is_success = True
|
|
352
|
+
mock_response.json.return_value = {
|
|
353
|
+
"data": {
|
|
354
|
+
"vulnerabilities": {
|
|
355
|
+
"nodes": [{"id": "1"}],
|
|
356
|
+
"pageInfo": {"hasNextPage": False, "endCursor": None},
|
|
357
|
+
}
|
|
358
|
+
}
|
|
359
|
+
}
|
|
360
|
+
|
|
361
|
+
progress_calls = []
|
|
362
|
+
|
|
363
|
+
def progress_callback(task_name, status, extra_data=None):
|
|
364
|
+
progress_calls.append((task_name, status, extra_data))
|
|
365
|
+
|
|
366
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
367
|
+
mock_client = AsyncMock()
|
|
368
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
369
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
370
|
+
|
|
371
|
+
result = await client.execute_paginated_query(
|
|
372
|
+
query="query { vulnerabilities }",
|
|
373
|
+
variables={"first": 100},
|
|
374
|
+
topic_key="vulnerabilities",
|
|
375
|
+
progress_callback=progress_callback,
|
|
376
|
+
task_name="Test Paginated Query",
|
|
377
|
+
)
|
|
378
|
+
|
|
379
|
+
assert len(result) == 1
|
|
380
|
+
assert any("fetched_page_1" in str(call) for call in progress_calls)
|
|
381
|
+
|
|
382
|
+
@pytest.mark.asyncio
|
|
383
|
+
async def test_execute_paginated_query_with_null_nodes(self):
|
|
384
|
+
"""Test paginated query when nodes is None"""
|
|
385
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
386
|
+
|
|
387
|
+
mock_response = Mock()
|
|
388
|
+
mock_response.is_success = True
|
|
389
|
+
mock_response.json.return_value = {
|
|
390
|
+
"data": {"vulnerabilities": {"nodes": None, "pageInfo": {"hasNextPage": False, "endCursor": None}}}
|
|
391
|
+
}
|
|
392
|
+
|
|
393
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
394
|
+
mock_client = AsyncMock()
|
|
395
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
396
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
397
|
+
|
|
398
|
+
result = await client.execute_paginated_query(
|
|
399
|
+
query="query { vulnerabilities }",
|
|
400
|
+
variables={"first": 100},
|
|
401
|
+
topic_key="vulnerabilities",
|
|
402
|
+
)
|
|
403
|
+
|
|
404
|
+
assert result == []
|
|
405
|
+
|
|
406
|
+
@pytest.mark.asyncio
|
|
407
|
+
async def test_execute_paginated_query_error_on_page(self):
|
|
408
|
+
"""Test paginated query with error on a page"""
|
|
409
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
410
|
+
|
|
411
|
+
# First page success
|
|
412
|
+
mock_response_1 = Mock()
|
|
413
|
+
mock_response_1.is_success = True
|
|
414
|
+
mock_response_1.json.return_value = {
|
|
415
|
+
"data": {
|
|
416
|
+
"vulnerabilities": {
|
|
417
|
+
"nodes": [{"id": "1"}],
|
|
418
|
+
"pageInfo": {"hasNextPage": True, "endCursor": "cursor1"},
|
|
419
|
+
}
|
|
420
|
+
}
|
|
421
|
+
}
|
|
422
|
+
|
|
423
|
+
# Second page fails
|
|
424
|
+
mock_response_2 = Mock()
|
|
425
|
+
mock_response_2.is_success = False
|
|
426
|
+
mock_response_2.status_code = 500
|
|
427
|
+
mock_response_2.text = "Server Error"
|
|
428
|
+
|
|
429
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
430
|
+
mock_client = AsyncMock()
|
|
431
|
+
mock_client.post = AsyncMock(side_effect=[mock_response_1, mock_response_2])
|
|
432
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
433
|
+
|
|
434
|
+
with patch(f"{PATH}.error_and_exit") as mock_error_exit:
|
|
435
|
+
mock_error_exit.side_effect = SystemExit(1)
|
|
436
|
+
|
|
437
|
+
# Should return partial results before error
|
|
438
|
+
with pytest.raises(SystemExit):
|
|
439
|
+
await client.execute_paginated_query(
|
|
440
|
+
query="query { vulnerabilities }",
|
|
441
|
+
variables={"first": 100},
|
|
442
|
+
topic_key="vulnerabilities",
|
|
443
|
+
)
|
|
444
|
+
|
|
445
|
+
@pytest.mark.asyncio
|
|
446
|
+
async def test_execute_paginated_query_empty_page_info(self):
|
|
447
|
+
"""Test paginated query with missing pageInfo"""
|
|
448
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
449
|
+
|
|
450
|
+
mock_response = Mock()
|
|
451
|
+
mock_response.is_success = True
|
|
452
|
+
mock_response.json.return_value = {"data": {"vulnerabilities": {"nodes": [{"id": "1"}], "pageInfo": {}}}}
|
|
453
|
+
|
|
454
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
455
|
+
mock_client = AsyncMock()
|
|
456
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
457
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
458
|
+
|
|
459
|
+
result = await client.execute_paginated_query(
|
|
460
|
+
query="query { vulnerabilities }",
|
|
461
|
+
variables={"first": 100},
|
|
462
|
+
topic_key="vulnerabilities",
|
|
463
|
+
)
|
|
464
|
+
|
|
465
|
+
assert len(result) == 1
|
|
466
|
+
assert result[0]["id"] == "1"
|
|
467
|
+
|
|
468
|
+
|
|
469
|
+
class TestAsyncWizGraphQLClientProgressCallback:
|
|
470
|
+
"""Test AsyncWizGraphQLClient._create_progress_callback method"""
|
|
471
|
+
|
|
472
|
+
def test_create_progress_callback_starting(self):
|
|
473
|
+
"""Test progress callback for starting status"""
|
|
474
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
475
|
+
|
|
476
|
+
mock_tracker = MagicMock()
|
|
477
|
+
callback = client._create_progress_callback(mock_tracker, "task_1", "vulnerabilities")
|
|
478
|
+
|
|
479
|
+
callback("test_task", "starting")
|
|
480
|
+
|
|
481
|
+
mock_tracker.update.assert_called_once()
|
|
482
|
+
call_args = mock_tracker.update.call_args
|
|
483
|
+
assert "Starting vulnerabilities" in call_args[1]["description"]
|
|
484
|
+
|
|
485
|
+
def test_create_progress_callback_requesting(self):
|
|
486
|
+
"""Test progress callback for requesting status"""
|
|
487
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
488
|
+
|
|
489
|
+
mock_tracker = MagicMock()
|
|
490
|
+
callback = client._create_progress_callback(mock_tracker, "task_1", "vulnerabilities")
|
|
491
|
+
|
|
492
|
+
callback("test_task", "requesting")
|
|
493
|
+
|
|
494
|
+
mock_tracker.update.assert_called_once()
|
|
495
|
+
call_args = mock_tracker.update.call_args
|
|
496
|
+
assert "Querying vulnerabilities" in call_args[1]["description"]
|
|
497
|
+
|
|
498
|
+
def test_create_progress_callback_completed(self):
|
|
499
|
+
"""Test progress callback for completed status"""
|
|
500
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
501
|
+
|
|
502
|
+
mock_tracker = MagicMock()
|
|
503
|
+
callback = client._create_progress_callback(mock_tracker, "task_1", "vulnerabilities")
|
|
504
|
+
|
|
505
|
+
callback("test_task", "completed")
|
|
506
|
+
|
|
507
|
+
mock_tracker.update.assert_called_once()
|
|
508
|
+
call_args = mock_tracker.update.call_args
|
|
509
|
+
assert "Completed vulnerabilities" in call_args[1]["description"]
|
|
510
|
+
|
|
511
|
+
def test_create_progress_callback_failed(self):
|
|
512
|
+
"""Test progress callback for failed status"""
|
|
513
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
514
|
+
|
|
515
|
+
mock_tracker = MagicMock()
|
|
516
|
+
callback = client._create_progress_callback(mock_tracker, "task_1", "vulnerabilities")
|
|
517
|
+
|
|
518
|
+
callback("test_task", "failed")
|
|
519
|
+
|
|
520
|
+
mock_tracker.update.assert_called_once()
|
|
521
|
+
call_args = mock_tracker.update.call_args
|
|
522
|
+
assert "Failed vulnerabilities" in call_args[1]["description"]
|
|
523
|
+
|
|
524
|
+
def test_create_progress_callback_fetched_page(self):
|
|
525
|
+
"""Test progress callback for fetched_page status"""
|
|
526
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
527
|
+
|
|
528
|
+
mock_tracker = MagicMock()
|
|
529
|
+
callback = client._create_progress_callback(mock_tracker, "task_1", "vulnerabilities")
|
|
530
|
+
|
|
531
|
+
callback("test_task", "fetched_page_1", {"nodes_count": 10, "total_nodes": 10})
|
|
532
|
+
|
|
533
|
+
mock_tracker.update.assert_called_once()
|
|
534
|
+
call_args = mock_tracker.update.call_args
|
|
535
|
+
assert "10 nodes fetched" in call_args[1]["description"]
|
|
536
|
+
|
|
537
|
+
def test_create_progress_callback_unknown_status(self):
|
|
538
|
+
"""Test progress callback with unknown status"""
|
|
539
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
540
|
+
|
|
541
|
+
mock_tracker = MagicMock()
|
|
542
|
+
callback = client._create_progress_callback(mock_tracker, "task_1", "vulnerabilities")
|
|
543
|
+
|
|
544
|
+
callback("test_task", "unknown_status")
|
|
545
|
+
|
|
546
|
+
# Should not crash, just not update
|
|
547
|
+
assert mock_tracker.update.call_count == 0
|
|
548
|
+
|
|
549
|
+
|
|
550
|
+
class TestAsyncWizGraphQLClientExecuteSingleQueryConfig:
|
|
551
|
+
"""Test AsyncWizGraphQLClient._execute_single_query_config method"""
|
|
552
|
+
|
|
553
|
+
@pytest.mark.asyncio
|
|
554
|
+
async def test_execute_single_query_config_success(self):
|
|
555
|
+
"""Test successful single query config execution"""
|
|
556
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
557
|
+
|
|
558
|
+
config = {
|
|
559
|
+
"type": WizVulnerabilityType.VULNERABILITY,
|
|
560
|
+
"query": "query { test }",
|
|
561
|
+
"variables": {"first": 100},
|
|
562
|
+
"topic_key": "vulnerabilities",
|
|
563
|
+
}
|
|
564
|
+
|
|
565
|
+
mock_response = Mock()
|
|
566
|
+
mock_response.is_success = True
|
|
567
|
+
mock_response.json.return_value = {
|
|
568
|
+
"data": {"vulnerabilities": {"nodes": [{"id": "1"}], "pageInfo": {"hasNextPage": False}}}
|
|
569
|
+
}
|
|
570
|
+
|
|
571
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
572
|
+
mock_client = AsyncMock()
|
|
573
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
574
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
575
|
+
|
|
576
|
+
query_type, results, error = await client._execute_single_query_config(config)
|
|
577
|
+
|
|
578
|
+
assert query_type == "vulnerability"
|
|
579
|
+
assert len(results) == 1
|
|
580
|
+
assert error is None
|
|
581
|
+
|
|
582
|
+
@pytest.mark.asyncio
|
|
583
|
+
async def test_execute_single_query_config_with_progress_tracker(self):
|
|
584
|
+
"""Test single query config with progress tracker"""
|
|
585
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
586
|
+
|
|
587
|
+
config = {
|
|
588
|
+
"type": WizVulnerabilityType.VULNERABILITY,
|
|
589
|
+
"query": "query { test }",
|
|
590
|
+
"variables": {"first": 100},
|
|
591
|
+
"topic_key": "vulnerabilities",
|
|
592
|
+
}
|
|
593
|
+
|
|
594
|
+
mock_response = Mock()
|
|
595
|
+
mock_response.is_success = True
|
|
596
|
+
mock_response.json.return_value = {
|
|
597
|
+
"data": {"vulnerabilities": {"nodes": [{"id": "1"}], "pageInfo": {"hasNextPage": False}}}
|
|
598
|
+
}
|
|
599
|
+
|
|
600
|
+
mock_tracker = MagicMock()
|
|
601
|
+
mock_tracker.add_task.return_value = "task_1"
|
|
602
|
+
|
|
603
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
604
|
+
mock_client = AsyncMock()
|
|
605
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
606
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
607
|
+
|
|
608
|
+
query_type, results, error = await client._execute_single_query_config(config, mock_tracker)
|
|
609
|
+
|
|
610
|
+
assert query_type == "vulnerability"
|
|
611
|
+
assert len(results) == 1
|
|
612
|
+
assert error is None
|
|
613
|
+
mock_tracker.add_task.assert_called_once()
|
|
614
|
+
mock_tracker.update.assert_called()
|
|
615
|
+
|
|
616
|
+
@pytest.mark.asyncio
|
|
617
|
+
async def test_execute_single_query_config_with_exception(self):
|
|
618
|
+
"""Test single query config with exception"""
|
|
619
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
620
|
+
|
|
621
|
+
config = {
|
|
622
|
+
"type": WizVulnerabilityType.VULNERABILITY,
|
|
623
|
+
"query": "query { test }",
|
|
624
|
+
"variables": {"first": 100},
|
|
625
|
+
"topic_key": "vulnerabilities",
|
|
626
|
+
}
|
|
627
|
+
|
|
628
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
629
|
+
mock_client = AsyncMock()
|
|
630
|
+
mock_client.post = AsyncMock(side_effect=Exception("Test error"))
|
|
631
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
632
|
+
|
|
633
|
+
with patch(f"{PATH}.error_and_exit") as mock_error_exit:
|
|
634
|
+
mock_error_exit.side_effect = SystemExit(1)
|
|
635
|
+
|
|
636
|
+
with pytest.raises(SystemExit):
|
|
637
|
+
await client._execute_single_query_config(config)
|
|
638
|
+
|
|
639
|
+
|
|
640
|
+
class TestAsyncWizGraphQLClientProcessConcurrentResults:
|
|
641
|
+
"""Test AsyncWizGraphQLClient._process_concurrent_results method"""
|
|
642
|
+
|
|
643
|
+
def test_process_concurrent_results_all_success(self):
|
|
644
|
+
"""Test processing concurrent results with all successes"""
|
|
645
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
646
|
+
|
|
647
|
+
results = [
|
|
648
|
+
("vulnerability", [{"id": "1"}], None),
|
|
649
|
+
("configuration_finding", [{"id": "2"}], None),
|
|
650
|
+
]
|
|
651
|
+
|
|
652
|
+
query_configs = [
|
|
653
|
+
{"type": WizVulnerabilityType.VULNERABILITY},
|
|
654
|
+
{"type": WizVulnerabilityType.CONFIGURATION},
|
|
655
|
+
]
|
|
656
|
+
|
|
657
|
+
processed = client._process_concurrent_results(results, query_configs)
|
|
658
|
+
|
|
659
|
+
assert len(processed) == 2
|
|
660
|
+
assert processed[0] == ("vulnerability", [{"id": "1"}], None)
|
|
661
|
+
assert processed[1] == ("configuration_finding", [{"id": "2"}], None)
|
|
662
|
+
|
|
663
|
+
def test_process_concurrent_results_with_exception(self):
|
|
664
|
+
"""Test processing concurrent results with exception"""
|
|
665
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
666
|
+
|
|
667
|
+
test_exception = Exception("Test error")
|
|
668
|
+
results = [
|
|
669
|
+
("vulnerability", [{"id": "1"}], None),
|
|
670
|
+
test_exception,
|
|
671
|
+
]
|
|
672
|
+
|
|
673
|
+
query_configs = [
|
|
674
|
+
{"type": WizVulnerabilityType.VULNERABILITY},
|
|
675
|
+
{"type": WizVulnerabilityType.CONFIGURATION},
|
|
676
|
+
]
|
|
677
|
+
|
|
678
|
+
processed = client._process_concurrent_results(results, query_configs)
|
|
679
|
+
|
|
680
|
+
assert len(processed) == 2
|
|
681
|
+
assert processed[0] == ("vulnerability", [{"id": "1"}], None)
|
|
682
|
+
assert processed[1][0] == "configuration_finding"
|
|
683
|
+
assert processed[1][1] == []
|
|
684
|
+
assert isinstance(processed[1][2], Exception)
|
|
685
|
+
|
|
686
|
+
|
|
687
|
+
class TestAsyncWizGraphQLClientExecuteConcurrentQueries:
|
|
688
|
+
"""Test AsyncWizGraphQLClient.execute_concurrent_queries method"""
|
|
689
|
+
|
|
690
|
+
@pytest.mark.asyncio
|
|
691
|
+
async def test_execute_concurrent_queries_success(self):
|
|
692
|
+
"""Test concurrent queries execution"""
|
|
693
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
694
|
+
|
|
695
|
+
query_configs = [
|
|
696
|
+
{
|
|
697
|
+
"type": WizVulnerabilityType.VULNERABILITY,
|
|
698
|
+
"query": "query { vulnerabilities }",
|
|
699
|
+
"variables": {"first": 100},
|
|
700
|
+
"topic_key": "vulnerabilities",
|
|
701
|
+
},
|
|
702
|
+
{
|
|
703
|
+
"type": WizVulnerabilityType.CONFIGURATION,
|
|
704
|
+
"query": "query { configurationFindings }",
|
|
705
|
+
"variables": {"first": 100},
|
|
706
|
+
"topic_key": "configurationFindings",
|
|
707
|
+
},
|
|
708
|
+
]
|
|
709
|
+
|
|
710
|
+
mock_response_1 = Mock()
|
|
711
|
+
mock_response_1.is_success = True
|
|
712
|
+
mock_response_1.json.return_value = {
|
|
713
|
+
"data": {"vulnerabilities": {"nodes": [{"id": "1"}], "pageInfo": {"hasNextPage": False}}}
|
|
714
|
+
}
|
|
715
|
+
|
|
716
|
+
mock_response_2 = Mock()
|
|
717
|
+
mock_response_2.is_success = True
|
|
718
|
+
mock_response_2.json.return_value = {
|
|
719
|
+
"data": {"configurationFindings": {"nodes": [{"id": "2"}], "pageInfo": {"hasNextPage": False}}}
|
|
720
|
+
}
|
|
721
|
+
|
|
722
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
723
|
+
mock_client = AsyncMock()
|
|
724
|
+
mock_client.post = AsyncMock(side_effect=[mock_response_1, mock_response_2])
|
|
725
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
726
|
+
|
|
727
|
+
results = await client.execute_concurrent_queries(query_configs)
|
|
728
|
+
|
|
729
|
+
assert len(results) == 2
|
|
730
|
+
assert results[0][0] == "vulnerability"
|
|
731
|
+
assert len(results[0][1]) == 1
|
|
732
|
+
assert results[1][0] == "configuration_finding"
|
|
733
|
+
assert len(results[1][1]) == 1
|
|
734
|
+
|
|
735
|
+
@pytest.mark.asyncio
|
|
736
|
+
async def test_execute_concurrent_queries_with_progress_tracker(self):
|
|
737
|
+
"""Test concurrent queries with progress tracker"""
|
|
738
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
739
|
+
|
|
740
|
+
query_configs = [
|
|
741
|
+
{
|
|
742
|
+
"type": WizVulnerabilityType.VULNERABILITY,
|
|
743
|
+
"query": "query { vulnerabilities }",
|
|
744
|
+
"variables": {"first": 100},
|
|
745
|
+
"topic_key": "vulnerabilities",
|
|
746
|
+
}
|
|
747
|
+
]
|
|
748
|
+
|
|
749
|
+
mock_response = Mock()
|
|
750
|
+
mock_response.is_success = True
|
|
751
|
+
mock_response.json.return_value = {
|
|
752
|
+
"data": {"vulnerabilities": {"nodes": [{"id": "1"}], "pageInfo": {"hasNextPage": False}}}
|
|
753
|
+
}
|
|
754
|
+
|
|
755
|
+
mock_tracker = MagicMock()
|
|
756
|
+
mock_tracker.add_task.return_value = "task_1"
|
|
757
|
+
|
|
758
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
759
|
+
mock_client = AsyncMock()
|
|
760
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
761
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
762
|
+
|
|
763
|
+
results = await client.execute_concurrent_queries(query_configs, mock_tracker)
|
|
764
|
+
|
|
765
|
+
assert len(results) == 1
|
|
766
|
+
mock_tracker.add_task.assert_called_once()
|
|
767
|
+
|
|
768
|
+
|
|
769
|
+
class TestRunAsyncQueries:
|
|
770
|
+
"""Test run_async_queries function"""
|
|
771
|
+
|
|
772
|
+
def test_run_async_queries_success(self):
|
|
773
|
+
"""Test run_async_queries with successful execution"""
|
|
774
|
+
endpoint = "https://api.wiz.io/graphql"
|
|
775
|
+
headers = {"Authorization": "Bearer test-token"}
|
|
776
|
+
query_configs = [
|
|
777
|
+
{
|
|
778
|
+
"type": WizVulnerabilityType.VULNERABILITY,
|
|
779
|
+
"query": "query { vulnerabilities }",
|
|
780
|
+
"variables": {"first": 100},
|
|
781
|
+
"topic_key": "vulnerabilities",
|
|
782
|
+
}
|
|
783
|
+
]
|
|
784
|
+
|
|
785
|
+
mock_response = Mock()
|
|
786
|
+
mock_response.is_success = True
|
|
787
|
+
mock_response.json.return_value = {
|
|
788
|
+
"data": {"vulnerabilities": {"nodes": [{"id": "1"}], "pageInfo": {"hasNextPage": False}}}
|
|
789
|
+
}
|
|
790
|
+
|
|
791
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
792
|
+
mock_client = AsyncMock()
|
|
793
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
794
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
795
|
+
|
|
796
|
+
results = run_async_queries(endpoint, headers, query_configs)
|
|
797
|
+
|
|
798
|
+
assert len(results) == 1
|
|
799
|
+
assert results[0][0] == "vulnerability"
|
|
800
|
+
assert len(results[0][1]) == 1
|
|
801
|
+
|
|
802
|
+
def test_run_async_queries_with_custom_parameters(self):
|
|
803
|
+
"""Test run_async_queries with custom max_concurrent and timeout"""
|
|
804
|
+
endpoint = "https://api.wiz.io/graphql"
|
|
805
|
+
headers = {"Authorization": "Bearer test-token"}
|
|
806
|
+
query_configs = [
|
|
807
|
+
{
|
|
808
|
+
"type": WizVulnerabilityType.VULNERABILITY,
|
|
809
|
+
"query": "query { vulnerabilities }",
|
|
810
|
+
"variables": {"first": 100},
|
|
811
|
+
"topic_key": "vulnerabilities",
|
|
812
|
+
}
|
|
813
|
+
]
|
|
814
|
+
|
|
815
|
+
mock_response = Mock()
|
|
816
|
+
mock_response.is_success = True
|
|
817
|
+
mock_response.json.return_value = {
|
|
818
|
+
"data": {"vulnerabilities": {"nodes": [{"id": "1"}], "pageInfo": {"hasNextPage": False}}}
|
|
819
|
+
}
|
|
820
|
+
|
|
821
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
822
|
+
mock_client = AsyncMock()
|
|
823
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
824
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
825
|
+
|
|
826
|
+
results = run_async_queries(endpoint, headers, query_configs, max_concurrent=10, timeout=120)
|
|
827
|
+
|
|
828
|
+
assert len(results) == 1
|
|
829
|
+
|
|
830
|
+
def test_run_async_queries_with_progress_tracker(self):
|
|
831
|
+
"""Test run_async_queries with progress tracker"""
|
|
832
|
+
endpoint = "https://api.wiz.io/graphql"
|
|
833
|
+
headers = {"Authorization": "Bearer test-token"}
|
|
834
|
+
query_configs = [
|
|
835
|
+
{
|
|
836
|
+
"type": WizVulnerabilityType.VULNERABILITY,
|
|
837
|
+
"query": "query { vulnerabilities }",
|
|
838
|
+
"variables": {"first": 100},
|
|
839
|
+
"topic_key": "vulnerabilities",
|
|
840
|
+
}
|
|
841
|
+
]
|
|
842
|
+
|
|
843
|
+
mock_response = Mock()
|
|
844
|
+
mock_response.is_success = True
|
|
845
|
+
mock_response.json.return_value = {
|
|
846
|
+
"data": {"vulnerabilities": {"nodes": [{"id": "1"}], "pageInfo": {"hasNextPage": False}}}
|
|
847
|
+
}
|
|
848
|
+
|
|
849
|
+
mock_tracker = MagicMock()
|
|
850
|
+
mock_tracker.add_task.return_value = "task_1"
|
|
851
|
+
|
|
852
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
853
|
+
mock_client = AsyncMock()
|
|
854
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
855
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
856
|
+
|
|
857
|
+
results = run_async_queries(endpoint, headers, query_configs, progress_tracker=mock_tracker)
|
|
858
|
+
|
|
859
|
+
assert len(results) == 1
|
|
860
|
+
mock_tracker.add_task.assert_called_once()
|
|
861
|
+
|
|
862
|
+
|
|
863
|
+
class TestEdgeCases:
|
|
864
|
+
"""Test edge cases and boundary conditions"""
|
|
865
|
+
|
|
866
|
+
@pytest.mark.asyncio
|
|
867
|
+
async def test_concurrent_request_limiting(self):
|
|
868
|
+
"""Test that max_concurrent limits concurrent requests"""
|
|
869
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql", max_concurrent=2)
|
|
870
|
+
|
|
871
|
+
assert client.max_concurrent == 2
|
|
872
|
+
assert client._semaphore._value == 2
|
|
873
|
+
|
|
874
|
+
@pytest.mark.asyncio
|
|
875
|
+
async def test_timeout_configuration(self):
|
|
876
|
+
"""Test that timeout is properly configured"""
|
|
877
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql", timeout=90.0)
|
|
878
|
+
|
|
879
|
+
assert client.timeout == 90.0
|
|
880
|
+
|
|
881
|
+
mock_response = Mock()
|
|
882
|
+
mock_response.is_success = True
|
|
883
|
+
mock_response.json.return_value = {"data": {"result": "success"}}
|
|
884
|
+
|
|
885
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
886
|
+
mock_client = AsyncMock()
|
|
887
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
888
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
889
|
+
|
|
890
|
+
await client.execute_query(query="query { test }")
|
|
891
|
+
|
|
892
|
+
call_kwargs = mock_client_class.call_args[1]
|
|
893
|
+
assert call_kwargs["timeout"] == 90.0
|
|
894
|
+
|
|
895
|
+
@pytest.mark.asyncio
|
|
896
|
+
async def test_empty_query_configs(self):
|
|
897
|
+
"""Test execute_concurrent_queries with empty configs"""
|
|
898
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
899
|
+
|
|
900
|
+
results = await client.execute_concurrent_queries([])
|
|
901
|
+
|
|
902
|
+
assert results == []
|
|
903
|
+
|
|
904
|
+
@pytest.mark.asyncio
|
|
905
|
+
async def test_query_with_null_topic_key_data(self):
|
|
906
|
+
"""Test paginated query when topic_key data is missing"""
|
|
907
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
908
|
+
|
|
909
|
+
mock_response = Mock()
|
|
910
|
+
mock_response.is_success = True
|
|
911
|
+
mock_response.json.return_value = {"data": {}}
|
|
912
|
+
|
|
913
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
914
|
+
mock_client = AsyncMock()
|
|
915
|
+
mock_client.post = AsyncMock(return_value=mock_response)
|
|
916
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
917
|
+
|
|
918
|
+
result = await client.execute_paginated_query(
|
|
919
|
+
query="query { vulnerabilities }",
|
|
920
|
+
variables={"first": 100},
|
|
921
|
+
topic_key="vulnerabilities",
|
|
922
|
+
)
|
|
923
|
+
|
|
924
|
+
assert result == []
|
|
925
|
+
|
|
926
|
+
@pytest.mark.asyncio
|
|
927
|
+
async def test_large_page_count(self):
|
|
928
|
+
"""Test paginated query with many pages"""
|
|
929
|
+
client = AsyncWizGraphQLClient(endpoint="https://api.wiz.io/graphql")
|
|
930
|
+
|
|
931
|
+
responses = []
|
|
932
|
+
for i in range(10):
|
|
933
|
+
mock_response = Mock()
|
|
934
|
+
mock_response.is_success = True
|
|
935
|
+
mock_response.json.return_value = {
|
|
936
|
+
"data": {
|
|
937
|
+
"vulnerabilities": {
|
|
938
|
+
"nodes": [{"id": str(i)}],
|
|
939
|
+
"pageInfo": {"hasNextPage": i < 9, "endCursor": f"cursor{i}" if i < 9 else None},
|
|
940
|
+
}
|
|
941
|
+
}
|
|
942
|
+
}
|
|
943
|
+
responses.append(mock_response)
|
|
944
|
+
|
|
945
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
946
|
+
mock_client = AsyncMock()
|
|
947
|
+
mock_client.post = AsyncMock(side_effect=responses)
|
|
948
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
949
|
+
|
|
950
|
+
result = await client.execute_paginated_query(
|
|
951
|
+
query="query { vulnerabilities }",
|
|
952
|
+
variables={"first": 100},
|
|
953
|
+
topic_key="vulnerabilities",
|
|
954
|
+
)
|
|
955
|
+
|
|
956
|
+
assert len(result) == 10
|
|
957
|
+
assert mock_client.post.call_count == 10
|
|
958
|
+
|
|
959
|
+
|
|
960
|
+
class TestIntegration:
|
|
961
|
+
"""Integration tests for complete workflows"""
|
|
962
|
+
|
|
963
|
+
def test_full_query_workflow(self):
|
|
964
|
+
"""Test complete workflow from run_async_queries to results"""
|
|
965
|
+
endpoint = "https://api.wiz.io/graphql"
|
|
966
|
+
headers = {"Authorization": "Bearer test-token"}
|
|
967
|
+
|
|
968
|
+
# Create multiple query configs
|
|
969
|
+
query_configs = [
|
|
970
|
+
{
|
|
971
|
+
"type": WizVulnerabilityType.VULNERABILITY,
|
|
972
|
+
"query": "query { vulnerabilities }",
|
|
973
|
+
"variables": {"first": 100},
|
|
974
|
+
"topic_key": "vulnerabilities",
|
|
975
|
+
},
|
|
976
|
+
{
|
|
977
|
+
"type": WizVulnerabilityType.CONFIGURATION,
|
|
978
|
+
"query": "query { configurationFindings }",
|
|
979
|
+
"variables": {"first": 100},
|
|
980
|
+
"topic_key": "configurationFindings",
|
|
981
|
+
},
|
|
982
|
+
{
|
|
983
|
+
"type": WizVulnerabilityType.HOST_FINDING,
|
|
984
|
+
"query": "query { hostFindings }",
|
|
985
|
+
"variables": {"first": 100},
|
|
986
|
+
"topic_key": "hostFindings",
|
|
987
|
+
},
|
|
988
|
+
]
|
|
989
|
+
|
|
990
|
+
# Mock responses for all queries
|
|
991
|
+
mock_response_1 = Mock()
|
|
992
|
+
mock_response_1.is_success = True
|
|
993
|
+
mock_response_1.json.return_value = {
|
|
994
|
+
"data": {"vulnerabilities": {"nodes": [{"id": "v1"}, {"id": "v2"}], "pageInfo": {"hasNextPage": False}}}
|
|
995
|
+
}
|
|
996
|
+
|
|
997
|
+
mock_response_2 = Mock()
|
|
998
|
+
mock_response_2.is_success = True
|
|
999
|
+
mock_response_2.json.return_value = {
|
|
1000
|
+
"data": {"configurationFindings": {"nodes": [{"id": "c1"}], "pageInfo": {"hasNextPage": False}}}
|
|
1001
|
+
}
|
|
1002
|
+
|
|
1003
|
+
mock_response_3 = Mock()
|
|
1004
|
+
mock_response_3.is_success = True
|
|
1005
|
+
mock_response_3.json.return_value = {
|
|
1006
|
+
"data": {
|
|
1007
|
+
"hostFindings": {
|
|
1008
|
+
"nodes": [{"id": "h1"}, {"id": "h2"}, {"id": "h3"}],
|
|
1009
|
+
"pageInfo": {"hasNextPage": False},
|
|
1010
|
+
}
|
|
1011
|
+
}
|
|
1012
|
+
}
|
|
1013
|
+
|
|
1014
|
+
with patch("httpx.AsyncClient") as mock_client_class:
|
|
1015
|
+
mock_client = AsyncMock()
|
|
1016
|
+
mock_client.post = AsyncMock(side_effect=[mock_response_1, mock_response_2, mock_response_3])
|
|
1017
|
+
mock_client_class.return_value.__aenter__.return_value = mock_client
|
|
1018
|
+
|
|
1019
|
+
results = run_async_queries(endpoint, headers, query_configs)
|
|
1020
|
+
|
|
1021
|
+
# Verify results
|
|
1022
|
+
assert len(results) == 3
|
|
1023
|
+
|
|
1024
|
+
# Check vulnerabilities
|
|
1025
|
+
assert results[0][0] == "vulnerability"
|
|
1026
|
+
assert len(results[0][1]) == 2
|
|
1027
|
+
assert results[0][2] is None
|
|
1028
|
+
|
|
1029
|
+
# Check configuration findings
|
|
1030
|
+
assert results[1][0] == "configuration_finding"
|
|
1031
|
+
assert len(results[1][1]) == 1
|
|
1032
|
+
assert results[1][2] is None
|
|
1033
|
+
|
|
1034
|
+
# Check host findings
|
|
1035
|
+
assert results[2][0] == "host_finding"
|
|
1036
|
+
assert len(results[2][1]) == 3
|
|
1037
|
+
assert results[2][2] is None
|