From 5ea51230b46383879c1f409b9e0b08c52801b113 Mon Sep 17 00:00:00 2001
From: Alexey Yarmosh <yarmoshalexey@gmail.com>
Date: Tue, 29 Oct 2024 16:05:52 +0100
Subject: [PATCH] feat: send adoption status to the probe

---
 src/lib/ws/gateway.ts                                        | 4 +++-
 .../integration/measurement/probe-communication.test.ts      | 5 +++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/lib/ws/gateway.ts b/src/lib/ws/gateway.ts
index 51f37eeb..7eb5c076 100644
--- a/src/lib/ws/gateway.ts
+++ b/src/lib/ws/gateway.ts
@@ -6,7 +6,7 @@ import { handleStatusUpdate } from '../../probe/handler/status.js';
 import { handleDnsUpdate } from '../../probe/handler/dns.js';
 import { handleStatsReport } from '../../probe/handler/stats.js';
 import { scopedLogger } from '../logger.js';
-import { probeOverride, getWsServer, PROBES_NAMESPACE, ServerSocket } from './server.js';
+import { probeOverride, getWsServer, PROBES_NAMESPACE, ServerSocket, adoptedProbes } from './server.js';
 import { probeMetadata } from './middleware/probe-metadata.js';
 import { errorHandler } from './helper/error-handler.js';
 import { subscribeWithHandler } from './helper/subscribe-handler.js';
@@ -23,6 +23,7 @@ io
 	.on('connect', errorHandler(async (socket: ServerSocket) => {
 		const probe = socket.data.probe;
 		const location = probeOverride.getUpdatedLocation(probe);
+		const isAdopted = !!adoptedProbes.getByIp(probe.ipAddress);
 
 		socket.emit('api:connect:alt-ips-token', {
 			token: await getAltIpsClient().generateToken(socket),
@@ -31,6 +32,7 @@ io
 		});
 
 		socket.emit('api:connect:location', location);
+		socket.emit('api:connect:adoption', { isAdopted });
 		logger.info(`ws client ${socket.id} connected from ${location.city}, ${location.country} [${probe.ipAddress} - ${location.network}]`);
 
 		// Handlers
diff --git a/test/tests/integration/measurement/probe-communication.test.ts b/test/tests/integration/measurement/probe-communication.test.ts
index 883a4c93..075cf368 100644
--- a/test/tests/integration/measurement/probe-communication.test.ts
+++ b/test/tests/integration/measurement/probe-communication.test.ts
@@ -17,6 +17,7 @@ describe('Create measurement request', () => {
 
 	const sandbox = sinon.createSandbox();
 	const locationHandlerStub = sandbox.stub();
+	const adoptionHandlerStub = sandbox.stub();
 	const requestHandlerStub = sandbox.stub();
 	const cryptoRandomString = sandbox.stub().returns('measurementid');
 
@@ -34,6 +35,7 @@ describe('Create measurement request', () => {
 
 		probe = await addFakeProbe({
 			'api:connect:location': locationHandlerStub,
+			'api:connect:adoption': adoptionHandlerStub,
 			'probe:measurement:request': requestHandlerStub,
 		});
 	});
@@ -75,6 +77,9 @@ describe('Create measurement request', () => {
 			network: 'The Constant Company LLC',
 			normalizedNetwork: 'the constant company llc',
 		}]);
+
+		expect(adoptionHandlerStub.callCount).to.equal(1);
+		expect(adoptionHandlerStub.firstCall.args).to.deep.equal([{ isAdopted: false }]);
 	});
 
 	it('should send and handle proper events during measurement request', async () => {