@dxos/messaging 0.6.13 → 0.6.14-main.2b6a0f3

package/package.json

@@ -1,23 +1,37 @@
 {
   "name": "@dxos/messaging",
-  "version": "0.6.13",
+  "version": "0.6.14-main.2b6a0f3",
   "description": "Messaging",
   "homepage": "https://dxos.org",
   "bugs": "https://github.com/dxos/dxos/issues",
   "license": "MIT",
   "author": "DXOS.org",
+  "sideEffects": false,
   "exports": {
     ".": {
       "browser": "./dist/lib/browser/index.mjs",
       "node": {
-        "default": "./dist/lib/node/index.cjs"
+        "require": "./dist/lib/node/index.cjs",
+        "default": "./dist/lib/node-esm/index.mjs"
       },
       "types": "./dist/types/src/index.d.ts"
+    },
+    "./testing": {
+      "browser": "./dist/lib/browser/testing/index.mjs",
+      "node": {
+        "require": "./dist/lib/node/testing/index.cjs",
+        "default": "./dist/lib/node-esm/testing/index.mjs"
+      },
+      "types": "./dist/types/src/testing/index.d.ts"
     }
   },
   "types": "dist/types/src/index.d.ts",
   "typesVersions": {
-    "*": {}
+    "*": {
+      "testing": [
+        "dist/types/src/testing/index.d.ts"
+      ]
+    }
   },
   "files": [
     "dist",
@@ -25,28 +39,25 @@
     "README.md"
   ],
   "dependencies": {
-    "@bufbuild/protobuf": "^2.0.0",
     "isomorphic-ws": "^5.0.0",
     "ws": "^8.14.2",
-    "@dxos/async": "0.6.13",
-    "@dxos/codec-protobuf": "0.6.13",
-    "@dxos/edge-client": "0.6.13",
-    "@dxos/context": "0.6.13",
-    "@dxos/invariant": "0.6.13",
-    "@dxos/keys": "0.6.13",
-    "@dxos/log": "0.6.13",
-    "@dxos/node-std": "0.6.13",
-    "@dxos/protocols": "0.6.13",
-    "@dxos/tracing": "0.6.13",
-    "@dxos/rpc": "0.6.13",
-    "@dxos/util": "0.6.13"
+    "@dxos/async": "0.6.14-main.2b6a0f3",
+    "@dxos/codec-protobuf": "0.6.14-main.2b6a0f3",
+    "@dxos/context": "0.6.14-main.2b6a0f3",
+    "@dxos/edge-client": "0.6.14-main.2b6a0f3",
+    "@dxos/invariant": "0.6.14-main.2b6a0f3",
+    "@dxos/keys": "0.6.14-main.2b6a0f3",
+    "@dxos/log": "0.6.14-main.2b6a0f3",
+    "@dxos/node-std": "0.6.14-main.2b6a0f3",
+    "@dxos/rpc": "0.6.14-main.2b6a0f3",
+    "@dxos/tracing": "0.6.14-main.2b6a0f3",
+    "@dxos/util": "0.6.14-main.2b6a0f3",
+    "@dxos/protocols": "0.6.14-main.2b6a0f3"
   },
   "devDependencies": {
-    "@types/node": "^18.11.9",
-    "earljs": "~0.1.10",
-    "typescript": "^5.5.4",
-    "wait-for-expect": "^3.0.2",
-    "@dxos/signal": "0.6.13"
+    "@dxos/signal": "0.6.14-main.2b6a0f3",
+    "@dxos/keyring": "0.6.14-main.2b6a0f3",
+    "@dxos/test-utils": "0.6.14-main.2b6a0f3"
   },
   "publishConfig": {
     "access": "public"

package/src/messenger.blueprint-test.ts

@@ -2,25 +2,21 @@
 // Copyright 2024 DXOS.org
 //
 
-import { expect } from 'earljs';
+import { describe, expect, test, onTestFinished } from 'vitest';
 
 import { asyncTimeout, latch, sleep } from '@dxos/async';
-import { type PublicKey } from '@dxos/keys';
-import { describe, test, afterTest } from '@dxos/test';
 import { range } from '@dxos/util';
 
-import { type SignalManager, WebsocketSignalManager } from './signal-manager';
+import { WebsocketSignalManager } from './signal-manager';
 import { type Message } from './signal-methods';
 import { messageEqual, PAYLOAD_1, PAYLOAD_2, PAYLOAD_3, TestBuilder } from './testing';
 
-export const messengerTests = (
-  signalManagerFactory: (identityKey: PublicKey, deviceKey: PublicKey) => Promise<SignalManager>,
-) => {
+export const messengerTests = (signalManagerFactory: TestBuilder['createSignalManager']) => {
   test('Message between peers', async () => {
     const builder = new TestBuilder({
       signalManagerFactory,
     });
-    afterTest(() => builder.close());
+    onTestFinished(() => builder.close());
     const peer1 = await builder.createPeer();
     const peer2 = await builder.createPeer();
 
@@ -39,11 +35,11 @@ export const messengerTests = (
     await promise;
   });
 
-  test('Message 3 peers', async () => {
+  test('Message 3 peers', { timeout: 1_000 }, async () => {
     const builder = new TestBuilder({
       signalManagerFactory,
     });
-    afterTest(() => builder.close());
+    onTestFinished(() => builder.close());
     const peer1 = await builder.createPeer();
     const peer2 = await builder.createPeer();
     const peer3 = await builder.createPeer();
@@ -83,13 +79,13 @@ export const messengerTests = (
       await peer2.messenger.sendMessage(message);
       await asyncTimeout(promise, 1_000);
     }
-  }).timeout(1_000);
+  });
 
-  test('Message routing', async () => {
+  test('Message routing', { timeout: 4_000 }, async () => {
     const builder = new TestBuilder({
       signalManagerFactory,
     });
-    afterTest(() => builder.close());
+    onTestFinished(() => builder.close());
     const peer1 = await builder.createPeer();
     const peer2 = await builder.createPeer();
 
@@ -140,13 +136,13 @@ export const messengerTests = (
       expect(messageEqual(message, onMessage2[0])).toEqual(true);
       expect(onMessage3.length === 0);
     }
-  }).timeout(4_000);
+  });
 
-  test('Unsubscribe listener', async () => {
+  test.skip('Unsubscribe listener', { timeout: 1_000 }, async () => {
     const builder = new TestBuilder({
       signalManagerFactory,
     });
-    afterTest(() => builder.close());
+    onTestFinished(() => builder.close());
     const peer1 = await builder.createPeer();
     const peer2 = await builder.createPeer();
 
@@ -207,15 +203,13 @@ export const messengerTests = (
       expect(messages1.length).toEqual(2);
       expect(messages2.length).toEqual(1);
     }
-  })
-    .tag('flaky')
-    .timeout(1_000);
+  });
 
-  test('re-entrant message', async () => {
+  test('re-entrant message', { timeout: 1_000 }, async () => {
     const builder = new TestBuilder({
       signalManagerFactory,
     });
-    afterTest(() => builder.close());
+    onTestFinished(() => builder.close());
     const peer1 = await builder.createPeer();
     const peer2 = await builder.createPeer();
 
@@ -245,9 +239,9 @@ export const messengerTests = (
       await peer1.messenger.sendMessage(message);
       await asyncTimeout(receivePromise, 1_000);
     }
-  }).timeout(1_000);
+  });
 
-  describe('Reliability', () => {
+  describe('Reliability', { timeout: 5_000 }, () => {
     test('message with non reliable connection', async () => {
       // Simulate unreliable connection.
       // Only each 3rd message is sent.
@@ -264,7 +258,7 @@ export const messengerTests = (
         signalManagerFactory,
         messageDisruption: unreliableConnection,
       });
-      afterTest(() => builder.close());
+      onTestFinished(() => builder.close());
       const peer1 = await builder.createPeer();
       await peer1.open();
       const peer2 = await builder.createPeer();
@@ -286,9 +280,9 @@ export const messengerTests = (
 
       // expect to receive 3 messages.
       await receivePromise;
-    }).timeout(5_000);
+    });
 
-    test('ignoring doubled messages', async () => {
+    test('ignoring doubled messages', { timeout: 5_000 }, async () => {
       // Message got doubled going through signal network.
       const doublingMessage = (data: Message) => [data, data];
 
@@ -296,7 +290,7 @@ export const messengerTests = (
         signalManagerFactory,
         messageDisruption: doublingMessage,
       });
-      afterTest(() => builder.close());
+      onTestFinished(() => builder.close());
       const peer1 = await builder.createPeer();
       await peer1.open();
       const peer2 = await builder.createPeer();
@@ -317,10 +311,10 @@ export const messengerTests = (
       await asyncTimeout(promise(), 1000);
       expect(count).toEqual(1);
     });
-  }).timeout(5_000);
+  });
 
-  describe('load', () => {
-    test('many connections to KUBE', async () => {
+  describe.skip('load', () => {
+    test('many connections to KUBE', { timeout: 5_000 }, async () => {
       const builder = new TestBuilder({
         signalManagerFactory: async () =>
           new WebsocketSignalManager([{ server: 'wss://dev.kube.dxos.org/.well-known/dx/signal' }]),
@@ -339,8 +333,6 @@ export const messengerTests = (
       });
 
       await sleep(1000000);
-    })
-      .tag('stress')
-      .timeout(5_000);
+    });
   });
 };

package/src/{messenger.test.ts → messenger.node.test.ts}

@@ -2,14 +2,13 @@
 // Copyright 2022 DXOS.org
 //
 
+import { afterAll, onTestFinished, beforeAll, test, describe } from 'vitest';
+
 import { asyncTimeout } from '@dxos/async';
-import { EdgeClient } from '@dxos/edge-client';
-import { type PublicKey } from '@dxos/keys';
 import { runTestSignalServer, type SignalServerRunner } from '@dxos/signal';
-import { afterAll, afterTest, beforeAll, openAndClose, test, describe } from '@dxos/test';
 
 import { messengerTests } from './messenger.blueprint-test';
-import { EdgeSignalManager, WebsocketSignalManager } from './signal-manager';
+import { WebsocketSignalManager } from './signal-manager';
 import { type Message } from './signal-methods';
 import { PAYLOAD_1, TestBuilder } from './testing';
 
@@ -26,12 +25,12 @@ describe('Messenger with WebsocketSignalManager', () => {
 
   messengerTests(async () => new WebsocketSignalManager([{ server: broker.url() }]));
 
-  test('Message with broken signal server', async () => {
+  test('Message with broken signal server', { timeout: 1000 }, async () => {
     const builder = new TestBuilder({
       signalManagerFactory: async () =>
         new WebsocketSignalManager([{ server: 'ws://broken.kube.' }, { server: broker.url() }]),
     });
-    afterTest(() => builder.close());
+    onTestFinished(() => builder.close());
     const peer1 = await builder.createPeer();
     const peer2 = await builder.createPeer();
 
@@ -46,17 +45,5 @@ describe('Messenger with WebsocketSignalManager', () => {
       await peer1.messenger.sendMessage(message);
       await asyncTimeout(receivePromise, 1_000);
     }
-  }).timeout(1_000);
-});
-
-// TODO(mykola): Expects wrangler dev in edge repo to run. Skip to pass CI.
-describe.skip('Messenger with EdgeSignalManager', () => {
-  const edgeSignalFactory = async (identityKey: PublicKey, deviceKey: PublicKey) => {
-    const client = new EdgeClient(identityKey.toHex(), deviceKey.toHex(), { socketEndpoint: 'ws://localhost:8787' });
-    await openAndClose(client);
-
-    return new EdgeSignalManager({ edgeConnection: client });
-  };
-
-  messengerTests(edgeSignalFactory);
+  });
 });

package/src/messenger.ts

@@ -57,7 +57,7 @@ export class Messenger {
   private _closed = true;
   private readonly _retryDelay: number;
 
-  constructor({ signalManager, retryDelay = 300 }: MessengerOptions) {
+  constructor({ signalManager, retryDelay = 1000 }: MessengerOptions) {
     this._signalManager = signalManager;
     this._retryDelay = retryDelay;
 

package/src/signal-client/{signal-client.test.ts → signal-client.node.test.ts}

@@ -2,15 +2,14 @@
 // Copyright 2020 DXOS.org
 //
 
-// eslint-disable-next-line @typescript-eslint/no-var-requires
-import { expect } from 'earljs';
+import { afterAll, beforeAll, describe, expect, test, onTestFinished } from 'vitest';
 
 import { asyncTimeout, waitForCondition } from '@dxos/async';
 import { type Any, type TaggedType } from '@dxos/codec-protobuf';
 import { PublicKey } from '@dxos/keys';
+import { log } from '@dxos/log';
 import { type TYPES } from '@dxos/protocols/proto';
 import { runTestSignalServer, type SignalServerRunner } from '@dxos/signal';
-import { afterAll, beforeAll, describe, test, afterTest } from '@dxos/test';
 import { ComplexSet, range } from '@dxos/util';
 
 import { SignalClient } from './signal-client';
@@ -30,7 +29,9 @@ describe('SignalClient', () => {
   });
 
   afterAll(async () => {
+    log.info('begin stop');
     await broker1.stop();
+    log.info('end stop');
   });
 
   test('message between 2 clients', async () => {
@@ -79,7 +80,7 @@ describe('SignalClient', () => {
     const message = createMessage(peer2, peer1);
 
     {
-      const receivedMessage = peer1.waitForNextMessage();
+      const receivedMessage = peer1.waitForNextMessage({ timeout: 1_000 });
       await peer2.client.sendMessage(message);
       expect(await receivedMessage).toEqual(message);
     }
@@ -90,7 +91,7 @@ describe('SignalClient', () => {
     {
       const receivedMessage = peer1.waitForNextMessage({ timeout: 200 });
       await peer2.client.sendMessage(message);
-      await expect(receivedMessage).toBeRejected();
+      await expect(receivedMessage).rejects.toBeDefined();
     }
   });
 
@@ -138,7 +139,7 @@ describe('SignalClient', () => {
       });
 
       void client.open();
-      afterTest(async () => {
+      onTestFinished(async () => {
         await client.close();
       });
       return {
@@ -146,10 +147,10 @@ describe('SignalClient', () => {
         identityKey,
         client,
         peerInfo: { peerKey: peerKey.toHex(), identityKey: identityKey.toHex() },
-        waitForNextMessage: async () => {
+        waitForNextMessage: async ({ timeout = 5_000 } = {}) => {
           return asyncTimeout(
             client.onMessage.waitFor(() => true),
-            5000,
+            timeout,
           );
         },
         waitForPeer: (peerId: PublicKey) => waitForCondition({ condition: () => peers.has(peerId) }),

package/src/signal-client/{signal-rpc-client.test.ts → signal-rpc-client.node.test.ts}

@@ -2,7 +2,7 @@
 // Copyright 2022 DXOS.org
 //
 
-import { expect } from 'earljs';
+import { afterAll, beforeAll, describe, expect, onTestFinished, test } from 'vitest';
 
 import { type Any } from '@dxos/codec-protobuf';
 import { PublicKey } from '@dxos/keys';
@@ -10,7 +10,6 @@ import { log } from '@dxos/log';
 import { schema } from '@dxos/protocols/proto';
 import { type Message as SignalMessage, type SwarmEvent } from '@dxos/protocols/proto/dxos/mesh/signal';
 import { runTestSignalServer, type SignalServerRunner } from '@dxos/signal';
-import { afterAll, afterTest, beforeAll, describe, test } from '@dxos/test';
 
 import { SignalRPCClient } from './signal-rpc-client';
 
@@ -28,11 +27,11 @@ describe('SignalRPCClient', () => {
   // TODO(burdon): Convert to TestBuilder pattern.
   const setupClient = async () => {
     const client = new SignalRPCClient({ url: broker.url() });
-    afterTest(async () => await client.close());
+    onTestFinished(async () => await client.close());
     return client;
   };
 
-  test('signal between 2 peers', async () => {
+  test('signal between 2 peers', { timeout: 2_000 }, async () => {
     const client1 = await setupClient();
     const client2 = await setupClient();
 
@@ -65,14 +64,12 @@ describe('SignalRPCClient', () => {
       payload,
     });
 
-    expect((await received).author).toEqual(peerId2.asUint8Array());
-    expect((await received).payload).toBeAnObjectWith(payload);
+    expect((await received).author).toEqual(peerId2.asBuffer());
+    expect((await received).payload).toEqual(expect.objectContaining(payload));
     void stream1.close();
-  })
-    .timeout(2_000)
-    .retries(2);
+  });
 
-  test('join', async () => {
+  test('join', { timeout: 2_000 }, async () => {
     const client1 = await setupClient();
     const client2 = await setupClient();
 
@@ -101,7 +98,5 @@ describe('SignalRPCClient', () => {
     expect((await promise).peerAvailable?.peer).toEqual(peerId2.asBuffer());
     void stream1.close();
     void stream2.close();
-  })
-    .timeout(2_000)
-    .retries(2);
+  });
 });

package/src/signal-manager/edge-signal-manager.ts

@@ -2,8 +2,6 @@
 // Copyright 2024 DXOS.org
 //
 
-import { AnySchema } from '@bufbuild/protobuf/wkt';
-
 import { Event } from '@dxos/async';
 import { Resource } from '@dxos/context';
 import { type EdgeConnection, protocol } from '@dxos/edge-client';
@@ -11,11 +9,13 @@ import { invariant } from '@dxos/invariant';
 import { PublicKey } from '@dxos/keys';
 import { log } from '@dxos/log';
 import { EdgeService } from '@dxos/protocols';
+import { type buf, bufWkt } from '@dxos/protocols/buf';
 import {
   SwarmRequestSchema,
   SwarmRequest_Action as SwarmRequestAction,
   SwarmResponseSchema,
   type Message as EdgeMessage,
+  type PeerSchema,
 } from '@dxos/protocols/buf/dxos/edge/messenger_pb';
 import { ComplexMap, ComplexSet } from '@dxos/util';
 
@@ -63,6 +63,7 @@ export class EdgeSignalManager extends Resource implements SignalManager {
     await this._edgeConnection.send(
       protocol.createMessage(SwarmRequestSchema, {
         serviceId: EdgeService.SWARM_SERVICE_ID,
+        source: createMessageSource(topic, peer),
         payload: { action: SwarmRequestAction.JOIN, swarmKeys: [topic.toHex()] },
       }),
     );
@@ -73,6 +74,7 @@ export class EdgeSignalManager extends Resource implements SignalManager {
     await this._edgeConnection.send(
       protocol.createMessage(SwarmRequestSchema, {
         serviceId: EdgeService.SWARM_SERVICE_ID,
+        source: createMessageSource(topic, peer),
         payload: { action: SwarmRequestAction.LEAVE, swarmKeys: [topic.toHex()] },
       }),
     );
@@ -88,7 +90,7 @@ export class EdgeSignalManager extends Resource implements SignalManager {
     }
 
     await this._edgeConnection.send(
-      protocol.createMessage(AnySchema, {
+      protocol.createMessage(bufWkt.AnySchema, {
         serviceId: EdgeService.SIGNAL_SERVICE_ID,
         source: message.author,
         target: [message.recipient],
@@ -126,7 +128,7 @@ export class EdgeSignalManager extends Resource implements SignalManager {
       return;
     }
     const oldPeers = this._swarmPeers.get(topic)!;
-    const timestamp = new Date(Date.parse(message.timestamp));
+    const timestamp = message.timestamp ? new Date(Date.parse(message.timestamp)) : new Date();
     const newPeers = new ComplexSet<PeerInfo>(PeerInfoHash, payload.peers);
 
     // Emit new available peers in the swarm.
@@ -155,8 +157,8 @@ export class EdgeSignalManager extends Resource implements SignalManager {
   }
 
   private _processMessage(message: EdgeMessage) {
-    invariant(protocol.getPayloadType(message) === AnySchema.typeName, 'Wrong payload type');
-    const payload = protocol.getPayload(message, AnySchema);
+    invariant(protocol.getPayloadType(message) === bufWkt.AnySchema.typeName, 'Wrong payload type');
+    const payload = protocol.getPayload(message, bufWkt.AnySchema);
     invariant(message.source, 'source is missing');
     invariant(message.target, 'target is missing');
     invariant(message.target.length === 1, 'target should have exactly one item');
@@ -188,3 +190,11 @@ export class EdgeSignalManager extends Resource implements SignalManager {
     }
   }
 }
+
+const createMessageSource = (topic: PublicKey, peerInfo: PeerInfo): buf.MessageInitShape<typeof PeerSchema> => {
+  return {
+    swarmKey: topic.toHex(),
+    identityKey: peerInfo.identityKey,
+    peerKey: peerInfo.peerKey,
+  };
+};

package/src/signal-manager/{websocket-signal-manager.test.ts → websocket-signal-manager.node.test.ts}

@@ -2,37 +2,31 @@
 // Copyright 2023 DXOS.org
 //
 
+import { afterAll, beforeAll, describe, test } from 'vitest';
+
 import { asyncTimeout, sleep } from '@dxos/async';
 import { PublicKey } from '@dxos/keys';
 import { runTestSignalServer, type SignalServerRunner } from '@dxos/signal';
-import { afterAll, beforeAll, describe, test, openAndClose } from '@dxos/test';
+import { openAndClose } from '@dxos/test-utils';
 
 import { WebsocketSignalManager } from './websocket-signal-manager';
 import { createMessage, expectPeerAvailable, expectReceivedMessage } from '../testing';
 
-describe('WebSocketSignalManager', () => {
+describe.skip('WebSocketSignalManager', () => {
   let broker1: SignalServerRunner;
   let broker2: SignalServerRunner;
 
   beforeAll(async () => {
-    if (!mochaExecutor.tags.includes('e2e')) {
-      return;
-    }
-
     broker1 = await runTestSignalServer({ port: 5001 });
     broker2 = await runTestSignalServer({ port: 5002 });
   });
 
   afterAll(() => {
-    if (!mochaExecutor.tags.includes('e2e')) {
-      return;
-    }
-
     void broker1.stop();
     void broker2.stop();
   });
 
-  test('join swarm with two brokers', async () => {
+  test('join swarm with two brokers', { timeout: 1_000 }, async () => {
     const client1 = new WebsocketSignalManager([{ server: broker1.url() }, { server: broker2.url() }]);
     const client2 = new WebsocketSignalManager([{ server: broker1.url() }]);
     const client3 = new WebsocketSignalManager([{ server: broker2.url() }]);
@@ -50,12 +44,9 @@ describe('WebSocketSignalManager', () => {
     await client3.join({ topic, peer: { peerKey: peer3.toHex() } });
 
     await Promise.all([joined12, joined13, joined21, joined31]);
-  })
-    .timeout(1_000)
-    .retries(2)
-    .tag('e2e');
+  });
 
-  test('join single swarm with doubled brokers', async () => {
+  test('join single swarm with doubled brokers', { timeout: 1_000 }, async () => {
     const client1 = new WebsocketSignalManager([{ server: broker1.url() }, { server: broker2.url() }]);
     const client2 = new WebsocketSignalManager([{ server: broker1.url() }, { server: broker2.url() }]);
     await openAndClose(client1, client2);
@@ -78,12 +69,9 @@ describe('WebSocketSignalManager', () => {
     await client1.sendMessage(message);
 
     await asyncTimeout(received, 1_000);
-  })
-    .timeout(1_000)
-    .retries(2)
-    .tag('e2e');
+  });
 
-  test('works with one broken server', async () => {
+  test('works with one broken server', { timeout: 1_000 }, async () => {
     const client1 = new WebsocketSignalManager([{ server: 'ws://broken.server/signal' }, { server: broker1.url() }]);
     const client2 = new WebsocketSignalManager([{ server: 'ws://broken.server/signal' }, { server: broker1.url() }]);
     await openAndClose(client1, client2);
@@ -97,12 +85,9 @@ describe('WebSocketSignalManager', () => {
     await client2.join({ topic, peer: { peerKey: peer2.toHex() } });
 
     await Promise.all([joined12, joined21]);
-  })
-    .timeout(1_000)
-    .retries(2)
-    .tag('e2e');
+  });
 
-  test('join two swarms with a broken signal server', async () => {
+  test('join two swarms with a broken signal server', { timeout: 1_000 }, async () => {
     const client1 = new WebsocketSignalManager([{ server: 'ws://broken.server/signal' }, { server: broker1.url() }]);
     const client2 = new WebsocketSignalManager([{ server: 'ws://broken.server/signal' }, { server: broker1.url() }]);
     await openAndClose(client1, client2);
@@ -122,8 +107,5 @@ describe('WebSocketSignalManager', () => {
     await client1.join({ topic: topic2, peer: { peerKey: peer1.toHex() } });
     await client2.join({ topic: topic2, peer: { peerKey: peer2.toHex() } });
     await Promise.all([joined212, joined221]);
-  })
-    .timeout(1_000)
-    .retries(2)
-    .tag('e2e');
+  });
 });

package/src/testing/test-builder.ts

@@ -2,14 +2,12 @@
 // Copyright 2022 DXOS.org
 //
 
-import { type PublicKey } from '@dxos/keys';
-
 import { TestPeer } from './test-peer';
 import { type SignalManager, MemorySignalManager, MemorySignalManagerContext } from '../signal-manager';
 import { type Message } from '../signal-methods';
 
 export type TestBuilderOptions = {
-  signalManagerFactory?: (identityKey: PublicKey, deviceKey: PublicKey) => Promise<SignalManager>;
+  signalManagerFactory?: (peer: TestPeer) => Promise<SignalManager>;
   messageDisruption?: (msg: Message) => Message[];
 };
 
@@ -19,10 +17,9 @@ export class TestBuilder {
 
   constructor(public options: TestBuilderOptions) {}
 
-  async createSignalManager(identityKey: PublicKey, deviceKey: PublicKey) {
+  async createSignalManager(peer: TestPeer): Promise<SignalManager> {
     const signalManager =
-      (await this.options.signalManagerFactory?.(identityKey, deviceKey)) ??
-      new MemorySignalManager(this._signalContext);
+      (await this.options.signalManagerFactory?.(peer)) ?? new MemorySignalManager(this._signalContext);
 
     if (this.options.messageDisruption) {
       // Imitates signal network disruptions (e. g. message doubling, ).