From bf5fb7fca4b9255c42a514b6a1b48274b22d1204 Mon Sep 17 00:00:00 2001
From: willclarktech <willclarktech@users.noreply.github.com>
Date: Tue, 22 Dec 2020 16:59:29 +0000
Subject: [PATCH] stargate: Actually use custom type urls in AminoTypes

---
 packages/stargate/src/aminotypes.spec.ts | 26 ++++++++++++++++++++++++
 packages/stargate/src/aminotypes.ts      | 11 +++++++---
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/packages/stargate/src/aminotypes.spec.ts b/packages/stargate/src/aminotypes.spec.ts
index 3053370d..a678aa2d 100644
--- a/packages/stargate/src/aminotypes.spec.ts
+++ b/packages/stargate/src/aminotypes.spec.ts
@@ -8,6 +8,18 @@ describe("AminoTypes", () => {
       expect(msgType).toEqual("cosmos-sdk/MsgDelegate");
     });
 
+    it("works with custom type url", () => {
+      const msgType = new AminoTypes({ "/my.CustomType": "my-sdk/CustomType" }).toAmino("/my.CustomType");
+      expect(msgType).toEqual("my-sdk/CustomType");
+    });
+
+    it("works with overridden type url", () => {
+      const msgType = new AminoTypes({
+        "/cosmos.staking.v1beta1.MsgDelegate": "my-override/MsgDelegate",
+      }).toAmino("/cosmos.staking.v1beta1.MsgDelegate");
+      expect(msgType).toEqual("my-override/MsgDelegate");
+    });
+
     it("throws for unknown type url", () => {
       expect(() => new AminoTypes().toAmino("/xxx.Unknown")).toThrowError(
         /Type URL does not exist in the Amino message type register./i,
@@ -21,6 +33,20 @@ describe("AminoTypes", () => {
       expect(msgUrl).toEqual("/cosmos.staking.v1beta1.MsgDelegate");
     });
 
+    it("works with custom type url", () => {
+      const msgType = new AminoTypes({ "/my.CustomType": "my-sdk/CustomType" }).fromAmino(
+        "my-sdk/CustomType",
+      );
+      expect(msgType).toEqual("/my.CustomType");
+    });
+
+    it("works with overridden type url", () => {
+      const msgType = new AminoTypes({
+        "/my.OverrideType": "cosmos-sdk/MsgDelegate",
+      }).fromAmino("cosmos-sdk/MsgDelegate");
+      expect(msgType).toEqual("/my.OverrideType");
+    });
+
     it("throws for unknown type url", () => {
       expect(() => new AminoTypes().fromAmino("cosmos-sdk/MsgUnknown")).toThrowError(
         /Type does not exist in the Amino message type register./i,
diff --git a/packages/stargate/src/aminotypes.ts b/packages/stargate/src/aminotypes.ts
index 0fa3a378..d8d20707 100644
--- a/packages/stargate/src/aminotypes.ts
+++ b/packages/stargate/src/aminotypes.ts
@@ -27,11 +27,16 @@ export class AminoTypes {
   private readonly register: Record<string, string>;
 
   public constructor(additions: Record<string, string> = {}) {
-    this.register = { ...defaultTypes, ...additions };
+    const additionalAminoTypes = Object.values(additions);
+    const filteredDefaultTypes = Object.entries(defaultTypes).reduce(
+      (acc, [key, value]) => (additionalAminoTypes.includes(value) ? acc : { ...acc, [key]: value }),
+      {},
+    );
+    this.register = { ...filteredDefaultTypes, ...additions };
   }
 
   public toAmino(typeUrl: string): string {
-    const type = defaultTypes[typeUrl];
+    const type = this.register[typeUrl];
     if (!type) {
       throw new Error(
         "Type URL does not exist in the Amino message type register. " +
@@ -43,7 +48,7 @@ export class AminoTypes {
   }
 
   public fromAmino(type: string): string {
-    const [typeUrl] = Object.entries(defaultTypes).find(([_typeUrl, value]) => value === type) ?? [];
+    const [typeUrl] = Object.entries(this.register).find(([_typeUrl, value]) => value === type) ?? [];
     if (!typeUrl) {
       throw new Error(
         "Type does not exist in the Amino message type register. " +