feat: add tests for updateOutputWithDescriptor

Issue: BTC-1451
BitGo · Oct 24, 2024 · f4c5db2 · f4c5db2
1 parent e8b4e5e
commit f4c5db2
Show file tree

Hide file tree

Showing 3 changed files with 131 additions and 11 deletions.
diff --git a/packages/wasm-miniscript/test/psbt.ts b/packages/wasm-miniscript/test/psbt.ts
@@ -1,10 +1,10 @@
 import * as utxolib from "@bitgo/utxo-lib";
 import * as assert from "node:assert";
-import { getPsbtFixtures } from "./psbtFixtures";
+import { getPsbtFixtures, PsbtStage } from "./psbtFixtures";
 import { Descriptor, Psbt } from "../js";
 
 import { getDescriptorForScriptType } from "./descriptorUtil";
-import { toUtxoPsbt, toWrappedPsbt, updateInputWithDescriptor } from "./psbt.util";
+import { assertEqualPsbt, toUtxoPsbt, toWrappedPsbt, updateInputWithDescriptor } from "./psbt.util";
 
 const rootWalletKeys = new utxolib.bitgo.RootWalletKeys(utxolib.testutil.getKeyTriple("wasm"));
 
@@ -14,15 +14,26 @@ function assertEqualBuffer(a: Buffer | Uint8Array, b: Buffer | Uint8Array, messa
 
 const fixtures = getPsbtFixtures(rootWalletKeys);
 
+function getWasmDescriptor(
+  scriptType: utxolib.bitgo.outputScripts.ScriptType2Of3,
+  scope: "internal" | "external",
+) {
+  return Descriptor.fromString(
+    getDescriptorForScriptType(rootWalletKeys, scriptType, scope),
+    "derivable",
+  );
+}
+
 function describeUpdateInputWithDescriptor(
   psbt: utxolib.bitgo.UtxoPsbt,
   scriptType: utxolib.bitgo.outputScripts.ScriptType2Of3,
 ) {
-  const fullSignedFixture = fixtures.find(
-    (f) => f.scriptType === scriptType && f.stage === "fullsigned",
-  );
-  if (!fullSignedFixture) {
-    throw new Error("Could not find fullsigned fixture");
+  function getFixtureAtStage(stage: PsbtStage) {
+    const f = fixtures.find((f) => f.scriptType === scriptType && f.stage === stage);
+    if (!f) {
+      throw new Error(`Could not find fixture for scriptType ${scriptType} and stage ${stage}`);
+    }
+    return f;
   }
 
   const descriptorStr = getDescriptorForScriptType(rootWalletKeys, scriptType, "internal");
@@ -33,7 +44,9 @@ function describeUpdateInputWithDescriptor(
     it("should update the input with the descriptor", function () {
       const wrappedPsbt = toWrappedPsbt(psbt);
       wrappedPsbt.updateInputWithDescriptor(0, descriptor.atDerivationIndex(index));
+      wrappedPsbt.updateOutputWithDescriptor(0, descriptor.atDerivationIndex(index));
       const updatedPsbt = toUtxoPsbt(wrappedPsbt);
+      assertEqualPsbt(updatedPsbt, getFixtureAtStage("unsigned").psbt);
       updatedPsbt.signAllInputsHD(rootWalletKeys.triple[0]);
       updatedPsbt.signAllInputsHD(rootWalletKeys.triple[2]);
       const wrappedSignedPsbt = toWrappedPsbt(updatedPsbt);
@@ -43,7 +56,11 @@ function describeUpdateInputWithDescriptor(
       assertEqualBuffer(updatedPsbt.toBuffer(), wrappedSignedPsbt.serialize());
 
       assertEqualBuffer(
-        fullSignedFixture.psbt.clone().finalizeAllInputs().extractTransaction().toBuffer(),
+        getFixtureAtStage("fullsigned")
+          .psbt.clone()
+          .finalizeAllInputs()
+          .extractTransaction()
+          .toBuffer(),
         updatedPsbt.extractTransaction().toBuffer(),
       );
     });
@@ -58,7 +75,11 @@ function describeUpdateInputWithDescriptor(
       cloned.finalizeAllInputs();
 
       assertEqualBuffer(
-        fullSignedFixture.psbt.clone().finalizeAllInputs().extractTransaction().toBuffer(),
+        getFixtureAtStage("fullsigned")
+          .psbt.clone()
+          .finalizeAllInputs()
+          .extractTransaction()
+          .toBuffer(),
         cloned.extractTransaction().toBuffer(),
       );
     });

diff --git a/packages/wasm-miniscript/test/psbt.util.ts b/packages/wasm-miniscript/test/psbt.util.ts
@@ -1,3 +1,4 @@
+import * as assert from "node:assert";
 import * as utxolib from "@bitgo/utxo-lib";
 import { Descriptor, Psbt } from "../js";
 
@@ -37,6 +38,16 @@ export function updateInputWithDescriptor(
   psbt.data.inputs[inputIndex] = toUtxoPsbt(wrappedPsbt).data.inputs[inputIndex];
 }
 
+export function updateOutputWithDescriptor(
+  psbt: utxolib.Psbt,
+  outputIndex: number,
+  descriptor: Descriptor,
+) {
+  const wrappedPsbt = toWrappedPsbt(psbt);
+  wrappedPsbt.updateOutputWithDescriptor(outputIndex, descriptor);
+  psbt.data.outputs[outputIndex] = toUtxoPsbt(wrappedPsbt).data.outputs[outputIndex];
+}
+
 export function finalizePsbt(psbt: utxolib.Psbt) {
   const wrappedPsbt = toWrappedPsbt(psbt);
   wrappedPsbt.finalize();
@@ -45,3 +56,90 @@ export function finalizePsbt(psbt: utxolib.Psbt) {
     psbt.data.inputs[i] = unwrappedPsbt.data.inputs[i];
   }
 }
+
+function toEntries(k: string, v: unknown, path: (string | number)[]): [] | [[string, unknown]] {
+  if (matchPath(path, ["data", "inputs", any, "sighashType"])) {
+    return [];
+  }
+  if (matchPath(path.slice(-1), ["unknownKeyVals"])) {
+    if (Array.isArray(v) && v.length === 0) {
+      return [];
+    }
+  }
+  return [[k, toPlainObject(v, path)]];
+}
+
+const any = Symbol("any");
+
+function matchPath(path: (string | number)[], pattern: (string | number | symbol)[]) {
+  if (path.length !== pattern.length) {
+    return false;
+  }
+  for (let i = 0; i < path.length; i++) {
+    if (pattern[i] !== any && path[i] !== pattern[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
+function normalizeBip32Derivation(v: unknown) {
+  if (!Array.isArray(v)) {
+    throw new Error("Expected bip32Derivation to be an array");
+  }
+  return (
+    [...v] as {
+      masterFingerprint: Buffer;
+      path: string;
+    }[]
+  )
+    .map((e) => {
+      let { path } = e;
+      if (path.startsWith("m/")) {
+        path = path.slice(2);
+      }
+      return {
+        ...e,
+        path,
+      };
+    })
+    .sort((a, b) => a.masterFingerprint.toString().localeCompare(b.masterFingerprint.toString()));
+}
+
+function toPlainObject(v: unknown, path: (string | number)[]) {
+  // psbts have fun getters and other types of irregular properties that we mash into shape here
+  if (v === null || v === undefined) {
+    return v;
+  }
+  if (
+    matchPath(path, ["data", "inputs", any, "bip32Derivation"]) ||
+    matchPath(path, ["data", "outputs", any, "bip32Derivation"])
+  ) {
+    v = normalizeBip32Derivation(v);
+  }
+  switch (typeof v) {
+    case "number":
+    case "bigint":
+    case "string":
+    case "boolean":
+      return v;
+    case "object":
+      if (v instanceof Buffer || v instanceof Uint8Array) {
+        return v.toString("hex");
+      }
+      if (Array.isArray(v)) {
+        return v.map((v, i) => toPlainObject(v, [...path, i]));
+      }
+      return Object.fromEntries(
+        Object.entries(v)
+          .flatMap(([k, v]) => toEntries(k, v, [...path, k]))
+          .sort(([a], [b]) => a.localeCompare(b)),
+      );
+    default:
+      throw new Error(`Unsupported type: ${typeof v}`);
+  }
+}
+
+export function assertEqualPsbt(a: utxolib.Psbt, b: utxolib.Psbt) {
+  assert.deepStrictEqual(toPlainObject(a, []), toPlainObject(b, []));
+}
diff --git a/packages/wasm-miniscript/test/psbtFixtures.ts b/packages/wasm-miniscript/test/psbtFixtures.ts
@@ -1,7 +1,7 @@
 import * as utxolib from "@bitgo/utxo-lib";
 import { RootWalletKeys } from "@bitgo/utxo-lib/dist/src/bitgo";
 
-type PsbtStage = "bare" | "unsigned" | "halfsigned" | "fullsigned";
+export type PsbtStage = "bare" | "unsigned" | "halfsigned" | "fullsigned";
 
 export function toPsbtWithPrevOutOnly(psbt: utxolib.bitgo.UtxoPsbt) {
   const psbtCopy = utxolib.bitgo.UtxoPsbt.createPsbt({
@@ -43,7 +43,8 @@ function getPsbtWithScriptTypeAndStage(
     [
       {
         value: BigInt(1e8 - 1000),
-        scriptType: "p2sh",
+        scriptType,
+        isInternalAddress: true,
       },
     ],
     utxolib.networks.bitcoin,