Skip to content

Commit 79fa552

Browse files
feat: Add connection resilience with auto-reconnect and heartbeat
- Add UnityConnection class with exponential backoff reconnection - Add ConnectionState enum (Disconnected, Connecting, Connected, Reconnecting) - Add heartbeat/ping-pong mechanism to detect stale connections - Emit connection state change events for UI feedback - Refactor McpUnity to use new UnityConnection class - Add Jest test suite with 14 tests for connection handling 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent c441c00 commit 79fa552

File tree

3 files changed

+873
-136
lines changed

3 files changed

+873
-136
lines changed
Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,190 @@
1+
import { jest, describe, it, expect, beforeEach, afterEach } from '@jest/globals';
2+
3+
// Mock WebSocket before importing modules that use it
4+
jest.unstable_mockModule('ws', () => ({
5+
default: jest.fn().mockImplementation(() => ({
6+
readyState: 1,
7+
onopen: null,
8+
onclose: null,
9+
onerror: null,
10+
onmessage: null,
11+
send: jest.fn(),
12+
close: jest.fn(),
13+
terminate: jest.fn(),
14+
ping: jest.fn(),
15+
on: jest.fn(),
16+
removeAllListeners: jest.fn()
17+
})),
18+
WebSocket: {
19+
CONNECTING: 0,
20+
OPEN: 1,
21+
CLOSING: 2,
22+
CLOSED: 3
23+
}
24+
}));
25+
26+
// Dynamic imports after mocking
27+
const { UnityConnection, ConnectionState } = await import('../unity/unityConnection');
28+
const { Logger, LogLevel } = await import('../utils/logger');
29+
const { McpUnityError, ErrorType } = await import('../utils/errors');
30+
31+
// Type imports
32+
import type { ConnectionStateChange } from '../unity/unityConnection';
33+
34+
// Create a logger that doesn't output anything (for testing)
35+
const createTestLogger = () => {
36+
process.env.LOGGING = 'false';
37+
process.env.LOGGING_FILE = 'false';
38+
return new Logger('Test', LogLevel.ERROR);
39+
};
40+
41+
describe('UnityConnection', () => {
42+
let connection: InstanceType<typeof UnityConnection>;
43+
let testLogger: InstanceType<typeof Logger>;
44+
45+
beforeEach(() => {
46+
testLogger = createTestLogger();
47+
connection = new UnityConnection(testLogger, {
48+
host: 'localhost',
49+
port: 8090,
50+
requestTimeout: 5000,
51+
clientName: 'TestClient',
52+
minReconnectDelay: 100,
53+
maxReconnectDelay: 1000,
54+
heartbeatInterval: 0
55+
});
56+
});
57+
58+
afterEach(() => {
59+
connection.disconnect();
60+
jest.clearAllMocks();
61+
});
62+
63+
describe('Initial State', () => {
64+
it('should start in disconnected state', () => {
65+
expect(connection.connectionState).toBe(ConnectionState.Disconnected);
66+
});
67+
68+
it('should not be connected initially', () => {
69+
expect(connection.isConnected).toBe(false);
70+
});
71+
72+
it('should not be connecting initially', () => {
73+
expect(connection.isConnecting).toBe(false);
74+
});
75+
76+
it('should have -1 for timeSinceLastPong before any connection', () => {
77+
expect(connection.timeSinceLastPong).toBe(-1);
78+
});
79+
});
80+
81+
describe('State Change Events', () => {
82+
it('should emit stateChange event when connect is called', (done) => {
83+
let firstEvent = true;
84+
connection.on('stateChange', (change: ConnectionStateChange) => {
85+
// Only check the first state change event
86+
if (firstEvent && change.currentState === ConnectionState.Connecting) {
87+
firstEvent = false;
88+
expect(change.previousState).toBe(ConnectionState.Disconnected);
89+
expect(change.currentState).toBe(ConnectionState.Connecting);
90+
done();
91+
}
92+
});
93+
94+
connection.connect().catch(() => {});
95+
});
96+
97+
it('should include reason in state change', (done) => {
98+
let eventReceived = false;
99+
connection.on('stateChange', (change: ConnectionStateChange) => {
100+
if (!eventReceived && change.currentState === ConnectionState.Connecting) {
101+
eventReceived = true;
102+
expect(change.reason).toBeDefined();
103+
done();
104+
}
105+
});
106+
107+
connection.connect().catch(() => {});
108+
});
109+
});
110+
111+
describe('Configuration', () => {
112+
it('should update configuration dynamically', () => {
113+
connection.updateConfig({ heartbeatInterval: 60000 });
114+
expect(connection.connectionState).toBe(ConnectionState.Disconnected);
115+
});
116+
});
117+
118+
describe('getStats', () => {
119+
it('should return correct stats in initial state', () => {
120+
const stats = connection.getStats();
121+
expect(stats.state).toBe(ConnectionState.Disconnected);
122+
expect(stats.reconnectAttempt).toBe(0);
123+
expect(stats.timeSinceLastPong).toBe(-1);
124+
expect(stats.isAwaitingPong).toBe(false);
125+
});
126+
});
127+
128+
describe('Disconnect', () => {
129+
it('should set state to disconnected on manual disconnect', () => {
130+
connection.disconnect('Test disconnect');
131+
expect(connection.connectionState).toBe(ConnectionState.Disconnected);
132+
});
133+
134+
it('should emit stateChange event when disconnecting from connecting state', (done) => {
135+
// First start connecting, then disconnect
136+
connection.on('stateChange', (change: ConnectionStateChange) => {
137+
if (change.currentState === ConnectionState.Disconnected &&
138+
change.previousState !== ConnectionState.Disconnected) {
139+
done();
140+
}
141+
});
142+
143+
// Start connection then immediately disconnect
144+
connection.connect().catch(() => {});
145+
// Give time for the connecting state to be set
146+
setTimeout(() => {
147+
connection.disconnect('Test disconnect');
148+
}, 10);
149+
});
150+
});
151+
152+
describe('Send', () => {
153+
it('should throw error when not connected', () => {
154+
expect(() => connection.send('test')).toThrow(McpUnityError);
155+
});
156+
});
157+
158+
describe('forceReconnect', () => {
159+
it('should trigger connecting state', () => {
160+
connection.forceReconnect();
161+
expect(connection.isConnecting).toBe(true);
162+
});
163+
});
164+
});
165+
166+
describe('ConnectionState Enum', () => {
167+
it('should have correct values', () => {
168+
expect(ConnectionState.Disconnected).toBe('disconnected');
169+
expect(ConnectionState.Connecting).toBe('connecting');
170+
expect(ConnectionState.Connected).toBe('connected');
171+
expect(ConnectionState.Reconnecting).toBe('reconnecting');
172+
});
173+
});
174+
175+
describe('Exponential Backoff Configuration', () => {
176+
it('should accept backoff configuration', () => {
177+
const testLogger = createTestLogger();
178+
const connection = new UnityConnection(testLogger, {
179+
host: 'localhost',
180+
port: 8090,
181+
requestTimeout: 5000,
182+
minReconnectDelay: 1000,
183+
maxReconnectDelay: 30000,
184+
reconnectBackoffMultiplier: 2
185+
});
186+
187+
expect(connection.connectionState).toBe(ConnectionState.Disconnected);
188+
connection.disconnect();
189+
});
190+
});

0 commit comments

Comments
 (0)