From a48e7ef0afe4e44070a85867eb86bf039a7ae3c8 Mon Sep 17 00:00:00 2001
From: Simonas Karuzas <simonas@not.cat>
Date: Mon, 5 Oct 2020 14:17:54 +0300
Subject: [PATCH] fix: Throwing errors for non existing entities

---
 __tests__/shared/verifiableData.ts     | 17 +++++++++++++
 packages/daf-typeorm/src/data-store.ts | 34 ++++++++++++++++++--------
 2 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/__tests__/shared/verifiableData.ts b/__tests__/shared/verifiableData.ts
index 2f14a92d7..71e801924 100644
--- a/__tests__/shared/verifiableData.ts
+++ b/__tests__/shared/verifiableData.ts
@@ -99,5 +99,22 @@ export default (testContext: {
       const presentationCount = await agent.dataStoreORMGetVerifiablePresentationsCount()
       expect(allPresentations.length).toEqual(presentationCount)
     })
+
+    it('should throw error for non existing verifiable credential', async () => {
+      await expect(
+        agent.dataStoreGetVerifiableCredential({
+          hash: 'foobar',
+        }),
+      ).rejects.toThrow('Verifiable credential not found')
+    })
+
+    it('should throw error for non existing verifiable presentation', async () => {
+      await expect(
+        agent.dataStoreGetVerifiablePresentation({
+          hash: 'foobar',
+        }),
+      ).rejects.toThrow('Verifiable presentation not found')
+    })
+
   })
 }
diff --git a/packages/daf-typeorm/src/data-store.ts b/packages/daf-typeorm/src/data-store.ts
index 9267cfaf5..7482ade37 100644
--- a/packages/daf-typeorm/src/data-store.ts
+++ b/packages/daf-typeorm/src/data-store.ts
@@ -41,8 +41,14 @@ export class DataStore implements IAgentPlugin {
   }
 
   async dataStoreGetMessage(args: IDataStoreGetMessageArgs): Promise<IMessage> {
-    const messageEntity = await (await this.dbConnection).getRepository(Message).findOneOrFail(args.id)
-    return createMessage(messageEntity)
+    try {
+      const messageEntity = await (await this.dbConnection).getRepository(Message).findOneOrFail(args.id, {
+        relations: ['credentials', 'presentations']
+      })
+      return createMessage(messageEntity)
+    } catch (e) {
+      throw Error('Message not found')
+    }
   }
 
   async dataStoreSaveVerifiableCredential(args: IDataStoreSaveVerifiableCredentialArgs): Promise<string> {
@@ -55,10 +61,14 @@ export class DataStore implements IAgentPlugin {
   async dataStoreGetVerifiableCredential(
     args: IDataStoreGetVerifiableCredentialArgs,
   ): Promise<VerifiableCredential> {
-    const credentialEntity = await (await this.dbConnection)
-      .getRepository(Credential)
-      .findOneOrFail(args.hash)
-    return credentialEntity.raw
+    try {
+      const credentialEntity = await (await this.dbConnection)
+        .getRepository(Credential)
+        .findOneOrFail(args.hash)
+      return credentialEntity.raw
+    } catch (e) {
+      throw Error('Verifiable credential not found')
+    }
   }
 
   async dataStoreSaveVerifiablePresentation(args: IDataStoreSaveVerifiablePresentationArgs): Promise<string> {
@@ -71,9 +81,13 @@ export class DataStore implements IAgentPlugin {
   async dataStoreGetVerifiablePresentation(
     args: IDataStoreGetVerifiablePresentationArgs,
   ): Promise<VerifiablePresentation> {
-    const presentationEntity = await (await this.dbConnection)
-      .getRepository(Presentation)
-      .findOneOrFail(args.hash)
-    return presentationEntity.raw
+    try {
+      const presentationEntity = await (await this.dbConnection)
+        .getRepository(Presentation)
+        .findOneOrFail(args.hash)
+      return presentationEntity.raw
+    } catch (e) {
+      throw Error('Verifiable presentation not found')
+    }
   }
 }