Fix typings for impls (#704)

penumbra-zone · Mar 7, 2024 · 6c91054 · 6c91054
1 parent 71e2c5d
commit 6c91054
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 29 deletions.
diff --git a/apps/extension/src/impls.ts b/apps/extension/src/impls.ts
@@ -1,4 +1,4 @@
-import { createPromiseClient } from '@connectrpc/connect';
+import { ServiceImpl, createPromiseClient } from '@connectrpc/connect';
 import { createGrpcWebTransport } from '@connectrpc/connect-web';
 import { createProxyImpl } from '@penumbra-zone/transport-dom/proxy';
 import { rethrowImplErrors } from './utils/rethrow-impl-errors';
@@ -22,9 +22,12 @@ import { stakingImpl } from '@penumbra-zone/router/src/grpc/staking';
 import { viewImpl } from '@penumbra-zone/router/src/grpc/view-protocol-server';
 
 import { localExtStorage } from '@penumbra-zone/storage';
+import { ServiceType } from '@bufbuild/protobuf';
 const grpcEndpoint = await localExtStorage.get('grpcEndpoint');
 
-const penumbraProxies = [
+type RpcImplTuple<T extends ServiceType> = [T, Partial<ServiceImpl<T>>];
+
+const penumbraProxies: RpcImplTuple<ServiceType>[] = [
   AppService,
   CompactBlockService,
   DexService,
@@ -44,12 +47,10 @@ const penumbraProxies = [
     ] as const,
 );
 
-export const rpcImpls = [
+export const rpcImpls: RpcImplTuple<ServiceType>[] = [
   // rpc local implementations
-  // @ts-expect-error TODO: accept partial impl
   [CustodyService, rethrowImplErrors(CustodyService, custodyImpl)],
   [SctService, rethrowImplErrors(SctService, sctImpl)],
-  // @ts-expect-error TODO: accept partial impl
   [StakingService, rethrowImplErrors(StakingService, stakingImpl)],
   [ViewService, rethrowImplErrors(ViewService, viewImpl)],
   // rpc remote proxies

diff --git a/apps/extension/src/utils/rethrow-impl-errors.test.ts b/apps/extension/src/utils/rethrow-impl-errors.test.ts
@@ -90,7 +90,7 @@ describe('rethrowImplErrors()', () => {
     describe('when the request succeeds', () => {
       it('returns the response', async () => {
         await expect(
-          wrappedServiceImplementation.getFoos(new FooRequest(), mockHandlerContext),
+          wrappedServiceImplementation.getFoos!(new FooRequest(), mockHandlerContext),
         ).resolves.toEqual(['foo1', 'foo2']);
       });
     });
@@ -103,7 +103,7 @@ describe('rethrowImplErrors()', () => {
         });
 
         await expect(() =>
-          wrappedServiceImplementation.getFoos(new FooRequest(), mockHandlerContext),
+          wrappedServiceImplementation.getFoos!(new FooRequest(), mockHandlerContext),
         ).rejects.toThrow(error);
       });
     });
@@ -116,7 +116,7 @@ describe('rethrowImplErrors()', () => {
         });
 
         await expect(() =>
-          wrappedServiceImplementation.getFoos(new FooRequest(), mockHandlerContext),
+          wrappedServiceImplementation.getFoos!(new FooRequest(), mockHandlerContext),
         ).rejects.toThrow(new ConnectError('oh no!'));
       });
     });
@@ -128,7 +128,7 @@ describe('rethrowImplErrors()', () => {
         expect.assertions(3);
 
         let whichBar = 1;
-        for await (const item of wrappedServiceImplementation.getBars(
+        for await (const item of wrappedServiceImplementation.getBars!(
           new BarRequest(),
           mockHandlerContext,
         )) {
@@ -146,7 +146,7 @@ describe('rethrowImplErrors()', () => {
         });
 
         await expect(async () => {
-          for await (const item of wrappedServiceImplementation.getBars(
+          for await (const item of wrappedServiceImplementation.getBars!(
             new BarRequest(),
             mockHandlerContext,
           )) {
@@ -165,7 +165,7 @@ describe('rethrowImplErrors()', () => {
         });
 
         await expect(async () => {
-          for await (const item of wrappedServiceImplementation.getBars(
+          for await (const item of wrappedServiceImplementation.getBars!(
             new BarRequest(),
             mockHandlerContext,
           )) {

diff --git a/apps/extension/src/utils/rethrow-impl-errors.ts b/apps/extension/src/utils/rethrow-impl-errors.ts
@@ -4,7 +4,6 @@ import {
   BiDiStreamingImpl,
   ClientStreamingImpl,
   HandlerContext,
-  MethodImpl,
   ServerStreamingImpl,
   UnaryImpl,
 } from '@connectrpc/connect/dist/cjs/implementation';
@@ -73,20 +72,16 @@ const isServerStreamingMethodKind = (
 
 export const rethrowImplErrors = <T extends ServiceType>(
   serviceType: T,
-  serviceImpl: ServiceImpl<T>,
-): ServiceImpl<T> =>
-  Object.fromEntries(
-    Object.entries(serviceImpl).map(
-      ([methodName, methodImplementation]: [
-        string,
-        MethodImpl<(typeof serviceType.methods)[string]>,
-      ]) => [
-        methodName,
-        isServerStreamingMethodKind(serviceType.methods[methodName]!.kind, methodImplementation)
-          ? wrapServerStreamingImpl(methodImplementation)
-          : isUnaryMethodKind(serviceType.methods[methodName]!.kind, methodImplementation)
-            ? wrapUnaryImpl(methodImplementation)
-            : wrapUnhandledImpl(methodImplementation, serviceType.methods[methodName]!.kind),
-      ],
-    ),
-  ) as ServiceImpl<T>;
+  serviceImpl: Partial<ServiceImpl<T>>,
+): Partial<ServiceImpl<T>> => {
+  const entries = Object.entries(serviceImpl).map(([methodName, methodImplementation]) => {
+    const methodKind = serviceType.methods[methodName]!.kind;
+    const wrappedMethodImpl = isServerStreamingMethodKind(methodKind, methodImplementation!)
+      ? wrapServerStreamingImpl(methodImplementation)
+      : isUnaryMethodKind(methodKind, methodImplementation!)
+        ? wrapUnaryImpl(methodImplementation)
+        : wrapUnhandledImpl(methodImplementation!, methodKind);
+    return [methodName, wrappedMethodImpl];
+  });
+  return Object.fromEntries(entries) as Partial<ServiceImpl<T>>;
+};