netra-zen 1.1.2__tar.gz → 1.2.0__tar.gz
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.
- {netra_zen-1.1.2/netra_zen.egg-info → netra_zen-1.2.0}/PKG-INFO +1 -1
- {netra_zen-1.1.2 → netra_zen-1.2.0/netra_zen.egg-info}/PKG-INFO +1 -1
- {netra_zen-1.1.2 → netra_zen-1.2.0}/netra_zen.egg-info/SOURCES.txt +1 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/pyproject.toml +1 -1
- {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/agent_cli.py +108 -9
- {netra_zen-1.1.2 → netra_zen-1.2.0}/setup.py +1 -1
- netra_zen-1.2.0/tests/test_timing_fix.py +243 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/LICENSE.md +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/MANIFEST.in +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/README.md +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/agent_interface/__init__.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/agent_interface/base_agent.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/config_example.json +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/docs/CACHE_TOKENS_GUIDE.md +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/docs/Cost_allocation.md +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/docs/DOLLAR_BUDGET_USAGE_EXAMPLES.md +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/docs/EXAMPLES.md +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/docs/MODEL_COLUMN_GUIDE.md +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/docs/apex_integration_test_plan.md +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/docs/zen_agent_cli_parallel_plan.md +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/netra_zen.egg-info/dependency_links.txt +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/netra_zen.egg-info/entry_points.txt +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/netra_zen.egg-info/requires.txt +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/netra_zen.egg-info/top_level.txt +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/prebuilt_commands_example.json +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/requirements-dev.txt +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/requirements.txt +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/__init__.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/__main__.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/agent_logs.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/bump_version.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/demo_log_collection.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/embed_release_credentials.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/test_apex_telemetry_debug.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/scripts/verify_log_transmission.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/setup.cfg +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/__init__.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/htmlcov/status.json +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_agent_interface.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_agent_logs.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_apex_integration.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_apex_streaming_fix.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_apex_telemetry.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_apex_telemetry_live.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_apex_telemetry_mock.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_apex_telemetry_regression.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_cli_extensions.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_cli_integration.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_direct_command_execution.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_dollar_budget_enhancement.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_event_batching_issue.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_permission_fix_windows.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_pricing_engine.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_runner.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_thread_handshake.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_thread_id_fix.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_thread_management.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_thread_resolution.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_workspace_detection.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_zen_commands.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_zen_integration.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_zen_metrics.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/tests/test_zen_unit.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/token_budget/__init__.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/token_budget/budget_manager.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/token_budget/models.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/token_budget/visualization.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/token_transparency/__init__.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/token_transparency/claude_pricing_engine.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/zen/__init__.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/zen/__main__.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/zen/telemetry/__init__.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/zen/telemetry/apex_telemetry.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/zen/telemetry/embedded_credentials.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/zen/telemetry/manager.py +0 -0
- {netra_zen-1.1.2 → netra_zen-1.2.0}/zen_orchestrator.py +0 -0
@@ -2883,22 +2883,49 @@ class WebSocketClient:
|
|
2883
2883
|
self.connected = True
|
2884
2884
|
return True
|
2885
2885
|
else:
|
2886
|
-
# Handshake failed -
|
2886
|
+
# Handshake failed - server not ready to process messages
|
2887
2887
|
self.debug.debug_print(
|
2888
|
-
"WARNING:
|
2888
|
+
"WARNING: Server handshake not completed - server not ready",
|
2889
2889
|
DebugLevel.BASIC,
|
2890
2890
|
style="yellow"
|
2891
2891
|
)
|
2892
2892
|
self.debug.debug_print(
|
2893
|
-
"Server
|
2893
|
+
"Server still completing lifecycle phases (Initialize → Authenticate → Handshake → Prepare → Processing)",
|
2894
2894
|
DebugLevel.VERBOSE,
|
2895
2895
|
style="yellow"
|
2896
2896
|
)
|
2897
2897
|
|
2898
|
-
#
|
2899
|
-
|
2900
|
-
|
2901
|
-
|
2898
|
+
# Wait briefly for server to complete its phases and retry
|
2899
|
+
safe_console_print("⏳ Waiting for server to complete initialization...", style="yellow",
|
2900
|
+
json_mode=self.config.json_mode, ci_mode=self.config.ci_mode)
|
2901
|
+
await asyncio.sleep(3.0) # Give server time to reach PROCESSING phase
|
2902
|
+
|
2903
|
+
# Retry handshake once more
|
2904
|
+
self.debug.debug_print(
|
2905
|
+
"Retrying handshake after delay...",
|
2906
|
+
DebugLevel.VERBOSE,
|
2907
|
+
style="cyan"
|
2908
|
+
)
|
2909
|
+
handshake_success = await self._perform_handshake()
|
2910
|
+
if handshake_success:
|
2911
|
+
safe_console_print(f"✅ Connected with thread ID: {self.current_thread_id}", style="green",
|
2912
|
+
json_mode=self.config.json_mode, ci_mode=self.config.ci_mode)
|
2913
|
+
self.connected = True
|
2914
|
+
return True
|
2915
|
+
else:
|
2916
|
+
# Server still not ready - fail the connection
|
2917
|
+
self.debug.debug_print(
|
2918
|
+
"ERROR: Server not ready after retry",
|
2919
|
+
DebugLevel.BASIC,
|
2920
|
+
style="red"
|
2921
|
+
)
|
2922
|
+
safe_console_print("❌ Server not ready to process messages. Please try again.", style="red",
|
2923
|
+
json_mode=self.config.json_mode, ci_mode=self.config.ci_mode)
|
2924
|
+
|
2925
|
+
# Close WebSocket and fail gracefully
|
2926
|
+
if self.ws:
|
2927
|
+
await self.ws.close()
|
2928
|
+
return False
|
2902
2929
|
except Exception as e:
|
2903
2930
|
self.debug.log_connection_attempt(method_name, self.config.ws_url, success=False, error=str(e))
|
2904
2931
|
self.debug.debug_print(
|
@@ -3027,7 +3054,8 @@ class WebSocketClient:
|
|
3027
3054
|
|
3028
3055
|
# Server enters HANDSHAKING phase after authentication
|
3029
3056
|
# and proactively sends handshake_response
|
3030
|
-
|
3057
|
+
# Use configured timeout or default to 10 seconds
|
3058
|
+
handshake_timeout = self.handshake_timeout if hasattr(self, 'handshake_timeout') and self.handshake_timeout else 10.0
|
3031
3059
|
start_time = asyncio.get_event_loop().time()
|
3032
3060
|
|
3033
3061
|
try:
|
@@ -3274,14 +3302,64 @@ class WebSocketClient:
|
|
3274
3302
|
)
|
3275
3303
|
|
3276
3304
|
# CRITICAL: Send acknowledgment with the SAME thread_id
|
3305
|
+
# Per WebSocket Client Lifecycle Guide, use "handshake_acknowledged"
|
3277
3306
|
ack_message = {
|
3278
|
-
"type": "
|
3307
|
+
"type": "handshake_acknowledged",
|
3279
3308
|
"thread_id": backend_thread_id, # Echo back the same ID
|
3280
3309
|
"timestamp": datetime.now(timezone.utc).isoformat()
|
3281
3310
|
}
|
3282
3311
|
|
3283
3312
|
await self.ws.send(json.dumps(ack_message))
|
3284
3313
|
|
3314
|
+
# Per WebSocket Client Lifecycle Guide, wait for handshake_complete
|
3315
|
+
# and then add delay for server to enter Phase 5 (Processing)
|
3316
|
+
self.debug.debug_print(
|
3317
|
+
"Waiting for handshake_complete confirmation...",
|
3318
|
+
DebugLevel.VERBOSE,
|
3319
|
+
style="cyan"
|
3320
|
+
)
|
3321
|
+
|
3322
|
+
# Wait briefly for handshake_complete message
|
3323
|
+
try:
|
3324
|
+
complete_msg = await asyncio.wait_for(self.ws.recv(), timeout=2.0)
|
3325
|
+
complete_data = json.loads(complete_msg)
|
3326
|
+
if complete_data.get('type') == 'handshake_complete':
|
3327
|
+
self.debug.debug_print(
|
3328
|
+
"✅ Received handshake_complete - Server at Phase 4 (Ready)",
|
3329
|
+
DebugLevel.VERBOSE,
|
3330
|
+
style="green"
|
3331
|
+
)
|
3332
|
+
|
3333
|
+
# CRITICAL: Add delay for server to enter Phase 5 (Processing)
|
3334
|
+
# Per documentation, 500ms is recommended
|
3335
|
+
self.debug.debug_print(
|
3336
|
+
"Waiting 500ms for server to enter Phase 5 (Processing)...",
|
3337
|
+
DebugLevel.VERBOSE,
|
3338
|
+
style="cyan"
|
3339
|
+
)
|
3340
|
+
await asyncio.sleep(0.5)
|
3341
|
+
|
3342
|
+
self.debug.debug_print(
|
3343
|
+
"✅ Server should now be in Phase 5 (Processing) - ready for messages",
|
3344
|
+
DebugLevel.VERBOSE,
|
3345
|
+
style="green"
|
3346
|
+
)
|
3347
|
+
except asyncio.TimeoutError:
|
3348
|
+
# If no handshake_complete, still add a delay to be safe
|
3349
|
+
self.debug.debug_print(
|
3350
|
+
"No handshake_complete received, adding safety delay",
|
3351
|
+
DebugLevel.VERBOSE,
|
3352
|
+
style="yellow"
|
3353
|
+
)
|
3354
|
+
await asyncio.sleep(0.5)
|
3355
|
+
except Exception as e:
|
3356
|
+
self.debug.debug_print(
|
3357
|
+
f"Error waiting for handshake_complete: {e}",
|
3358
|
+
DebugLevel.VERBOSE,
|
3359
|
+
style="yellow"
|
3360
|
+
)
|
3361
|
+
await asyncio.sleep(0.5)
|
3362
|
+
|
3285
3363
|
return True
|
3286
3364
|
|
3287
3365
|
def _get_platform_cache_path(self) -> Path:
|
@@ -3622,6 +3700,27 @@ class WebSocketClient:
|
|
3622
3700
|
if not self.ws:
|
3623
3701
|
raise RuntimeError("WebSocket not connected")
|
3624
3702
|
|
3703
|
+
# Check if connection handshake is fully complete
|
3704
|
+
if not self.connected:
|
3705
|
+
self.debug.debug_print(
|
3706
|
+
"ERROR: Connection not ready - handshake not complete",
|
3707
|
+
DebugLevel.BASIC,
|
3708
|
+
style="red"
|
3709
|
+
)
|
3710
|
+
safe_console_print(
|
3711
|
+
"\n❌ ERROR: Cannot send message - server not ready",
|
3712
|
+
style="red",
|
3713
|
+
json_mode=self.config.json_mode,
|
3714
|
+
ci_mode=self.config.ci_mode
|
3715
|
+
)
|
3716
|
+
safe_console_print(
|
3717
|
+
"The server is still completing its initialization phases.",
|
3718
|
+
style="yellow",
|
3719
|
+
json_mode=self.config.json_mode,
|
3720
|
+
ci_mode=self.config.ci_mode
|
3721
|
+
)
|
3722
|
+
raise RuntimeError("Connection not ready - handshake incomplete. Server needs to reach Processing phase.")
|
3723
|
+
|
3625
3724
|
# Generate run_id
|
3626
3725
|
self.run_id = f"cli_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{os.getpid()}"
|
3627
3726
|
|
@@ -8,7 +8,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
|
|
8
8
|
|
9
9
|
setup(
|
10
10
|
name="netra-zen",
|
11
|
-
version="1.
|
11
|
+
version="1.2.0",
|
12
12
|
author=" Systems",
|
13
13
|
author_email="pypi@netrasystems.ai",
|
14
14
|
description="Multi-instance Claude orchestrator for parallel task execution",
|
@@ -0,0 +1,243 @@
|
|
1
|
+
"""
|
2
|
+
Test for the timing fix that ensures CLI waits for server handshake
|
3
|
+
before sending messages.
|
4
|
+
|
5
|
+
This test verifies that:
|
6
|
+
1. CLI waits for handshake_response before marking connection as ready
|
7
|
+
2. CLI retries handshake with delay if server isn't ready
|
8
|
+
3. CLI properly fails connection if server doesn't complete handshake
|
9
|
+
4. No messages are sent before server reaches PROCESSING phase
|
10
|
+
"""
|
11
|
+
|
12
|
+
import asyncio
|
13
|
+
import json
|
14
|
+
import unittest
|
15
|
+
from unittest.mock import AsyncMock, MagicMock, patch, call
|
16
|
+
from datetime import datetime
|
17
|
+
import sys
|
18
|
+
import os
|
19
|
+
|
20
|
+
# Add parent directory to path for imports
|
21
|
+
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
22
|
+
|
23
|
+
from scripts.agent_cli import WebSocketClient
|
24
|
+
from shared import Environment, EnvironmentConfig
|
25
|
+
|
26
|
+
|
27
|
+
class TestTimingFix(unittest.TestCase):
|
28
|
+
"""Test suite for CLI timing fix"""
|
29
|
+
|
30
|
+
def setUp(self):
|
31
|
+
"""Set up test fixtures"""
|
32
|
+
self.config = MagicMock()
|
33
|
+
self.config.ws_url = "ws://localhost:8000/ws"
|
34
|
+
self.config.environment = Environment.LOCAL
|
35
|
+
self.config.json_mode = False
|
36
|
+
self.config.ci_mode = False
|
37
|
+
self.token = "test_token"
|
38
|
+
self.debug = MagicMock()
|
39
|
+
|
40
|
+
@patch('scripts.agent_cli.websockets.connect')
|
41
|
+
async def test_successful_handshake_on_first_try(self, mock_connect):
|
42
|
+
"""Test successful handshake on first attempt"""
|
43
|
+
# Mock WebSocket connection
|
44
|
+
mock_ws = AsyncMock()
|
45
|
+
mock_connect.return_value = mock_ws
|
46
|
+
|
47
|
+
# Mock handshake response
|
48
|
+
handshake_response = {
|
49
|
+
"type": "handshake_response",
|
50
|
+
"thread_id": "test-thread-123",
|
51
|
+
"user_id": "user-456",
|
52
|
+
"session_id": "session-789"
|
53
|
+
}
|
54
|
+
mock_ws.recv.return_value = json.dumps(handshake_response)
|
55
|
+
|
56
|
+
# Create client and connect
|
57
|
+
client = WebSocketClient(
|
58
|
+
self.config, self.token, self.debug,
|
59
|
+
send_logs=False, logs_count=10,
|
60
|
+
logs_project=None, logs_path=None,
|
61
|
+
logs_user=None, handshake_timeout=5
|
62
|
+
)
|
63
|
+
|
64
|
+
result = await client.connect()
|
65
|
+
|
66
|
+
# Assertions
|
67
|
+
self.assertTrue(result)
|
68
|
+
self.assertTrue(client.connected)
|
69
|
+
self.assertEqual(client.current_thread_id, "test-thread-123")
|
70
|
+
mock_ws.close.assert_not_called()
|
71
|
+
|
72
|
+
@patch('scripts.agent_cli.websockets.connect')
|
73
|
+
@patch('scripts.agent_cli.asyncio.sleep')
|
74
|
+
async def test_handshake_retry_after_delay(self, mock_sleep, mock_connect):
|
75
|
+
"""Test handshake retry after delay when server not ready"""
|
76
|
+
# Mock WebSocket connection
|
77
|
+
mock_ws = AsyncMock()
|
78
|
+
mock_connect.return_value = mock_ws
|
79
|
+
|
80
|
+
# First attempt: timeout (server not ready)
|
81
|
+
# Second attempt: successful handshake
|
82
|
+
handshake_response = {
|
83
|
+
"type": "handshake_response",
|
84
|
+
"thread_id": "test-thread-123",
|
85
|
+
"user_id": "user-456",
|
86
|
+
"session_id": "session-789"
|
87
|
+
}
|
88
|
+
|
89
|
+
# Configure recv to timeout first, then succeed
|
90
|
+
mock_ws.recv.side_effect = [
|
91
|
+
asyncio.TimeoutError(), # First handshake attempt times out
|
92
|
+
json.dumps(handshake_response) # Second attempt succeeds
|
93
|
+
]
|
94
|
+
|
95
|
+
# Create client and connect
|
96
|
+
client = WebSocketClient(
|
97
|
+
self.config, self.token, self.debug,
|
98
|
+
send_logs=False, logs_count=10,
|
99
|
+
logs_project=None, logs_path=None,
|
100
|
+
logs_user=None, handshake_timeout=1
|
101
|
+
)
|
102
|
+
|
103
|
+
result = await client.connect()
|
104
|
+
|
105
|
+
# Assertions
|
106
|
+
self.assertTrue(result)
|
107
|
+
self.assertTrue(client.connected)
|
108
|
+
self.assertEqual(client.current_thread_id, "test-thread-123")
|
109
|
+
mock_sleep.assert_called_once_with(3.0) # Verify delay was applied
|
110
|
+
mock_ws.close.assert_not_called()
|
111
|
+
|
112
|
+
@patch('scripts.agent_cli.websockets.connect')
|
113
|
+
@patch('scripts.agent_cli.asyncio.sleep')
|
114
|
+
async def test_connection_fails_if_no_handshake(self, mock_sleep, mock_connect):
|
115
|
+
"""Test connection fails properly when server never completes handshake"""
|
116
|
+
# Mock WebSocket connection
|
117
|
+
mock_ws = AsyncMock()
|
118
|
+
mock_connect.return_value = mock_ws
|
119
|
+
|
120
|
+
# Both handshake attempts timeout
|
121
|
+
mock_ws.recv.side_effect = [
|
122
|
+
asyncio.TimeoutError(), # First attempt
|
123
|
+
asyncio.TimeoutError() # Second attempt after retry
|
124
|
+
]
|
125
|
+
|
126
|
+
# Create client and connect
|
127
|
+
client = WebSocketClient(
|
128
|
+
self.config, self.token, self.debug,
|
129
|
+
send_logs=False, logs_count=10,
|
130
|
+
logs_project=None, logs_path=None,
|
131
|
+
logs_user=None, handshake_timeout=1
|
132
|
+
)
|
133
|
+
|
134
|
+
result = await client.connect()
|
135
|
+
|
136
|
+
# Assertions
|
137
|
+
self.assertFalse(result) # Connection should fail
|
138
|
+
self.assertFalse(client.connected) # Should not be marked as connected
|
139
|
+
self.assertIsNone(client.current_thread_id) # No thread ID assigned
|
140
|
+
mock_sleep.assert_called_once_with(3.0) # Verify retry delay
|
141
|
+
mock_ws.close.assert_called_once() # WebSocket should be closed
|
142
|
+
|
143
|
+
@patch('scripts.agent_cli.websockets.connect')
|
144
|
+
async def test_no_message_sent_without_handshake(self, mock_connect):
|
145
|
+
"""Test that send_message fails if handshake not completed"""
|
146
|
+
# Mock WebSocket connection
|
147
|
+
mock_ws = AsyncMock()
|
148
|
+
mock_connect.return_value = mock_ws
|
149
|
+
|
150
|
+
# Handshake times out - no response
|
151
|
+
mock_ws.recv.side_effect = asyncio.TimeoutError()
|
152
|
+
|
153
|
+
# Create client and try to connect
|
154
|
+
client = WebSocketClient(
|
155
|
+
self.config, self.token, self.debug,
|
156
|
+
send_logs=False, logs_count=10,
|
157
|
+
logs_project=None, logs_path=None,
|
158
|
+
logs_user=None, handshake_timeout=1
|
159
|
+
)
|
160
|
+
|
161
|
+
# Connection should fail due to no handshake
|
162
|
+
result = await client.connect()
|
163
|
+
self.assertFalse(result)
|
164
|
+
|
165
|
+
# Attempting to send message should fail
|
166
|
+
with self.assertRaises(Exception) as context:
|
167
|
+
await client.send_message("test message")
|
168
|
+
|
169
|
+
# Verify no message was sent via WebSocket
|
170
|
+
mock_ws.send.assert_not_called()
|
171
|
+
|
172
|
+
@patch('scripts.agent_cli.websockets.connect')
|
173
|
+
async def test_connection_established_is_not_handshake(self, mock_connect):
|
174
|
+
"""Test that connection_established event is not treated as handshake"""
|
175
|
+
# Mock WebSocket connection
|
176
|
+
mock_ws = AsyncMock()
|
177
|
+
mock_connect.return_value = mock_ws
|
178
|
+
|
179
|
+
# Send connection_established instead of handshake_response
|
180
|
+
connection_event = {
|
181
|
+
"type": "connection_established",
|
182
|
+
"message": "Connected to server"
|
183
|
+
}
|
184
|
+
|
185
|
+
# Then timeout (no actual handshake)
|
186
|
+
mock_ws.recv.side_effect = [
|
187
|
+
json.dumps(connection_event),
|
188
|
+
asyncio.TimeoutError()
|
189
|
+
]
|
190
|
+
|
191
|
+
# Create client and connect
|
192
|
+
client = WebSocketClient(
|
193
|
+
self.config, self.token, self.debug,
|
194
|
+
send_logs=False, logs_count=10,
|
195
|
+
logs_project=None, logs_path=None,
|
196
|
+
logs_user=None, handshake_timeout=1
|
197
|
+
)
|
198
|
+
|
199
|
+
result = await client.connect()
|
200
|
+
|
201
|
+
# Connection should fail - connection_established is not a handshake
|
202
|
+
self.assertFalse(result)
|
203
|
+
self.assertFalse(client.connected)
|
204
|
+
self.assertIsNone(client.current_thread_id)
|
205
|
+
|
206
|
+
|
207
|
+
def run_async_test(coro):
|
208
|
+
"""Helper to run async tests"""
|
209
|
+
loop = asyncio.get_event_loop()
|
210
|
+
return loop.run_until_complete(coro)
|
211
|
+
|
212
|
+
|
213
|
+
class AsyncTestRunner(unittest.TestCase):
|
214
|
+
"""Wrapper to run async tests"""
|
215
|
+
|
216
|
+
def test_successful_handshake(self):
|
217
|
+
test = TestTimingFix()
|
218
|
+
test.setUp()
|
219
|
+
run_async_test(test.test_successful_handshake_on_first_try())
|
220
|
+
|
221
|
+
def test_retry_after_delay(self):
|
222
|
+
test = TestTimingFix()
|
223
|
+
test.setUp()
|
224
|
+
run_async_test(test.test_handshake_retry_after_delay())
|
225
|
+
|
226
|
+
def test_connection_fails(self):
|
227
|
+
test = TestTimingFix()
|
228
|
+
test.setUp()
|
229
|
+
run_async_test(test.test_connection_fails_if_no_handshake())
|
230
|
+
|
231
|
+
def test_no_early_messages(self):
|
232
|
+
test = TestTimingFix()
|
233
|
+
test.setUp()
|
234
|
+
run_async_test(test.test_no_message_sent_without_handshake())
|
235
|
+
|
236
|
+
def test_connection_established_ignored(self):
|
237
|
+
test = TestTimingFix()
|
238
|
+
test.setUp()
|
239
|
+
run_async_test(test.test_connection_established_is_not_handshake())
|
240
|
+
|
241
|
+
|
242
|
+
if __name__ == '__main__':
|
243
|
+
unittest.main()
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|