From 67de2133744e90aade2c93d9401b844ce35256b4 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 24 Mar 2023 15:24:30 +1100
Subject: [PATCH] Account for data key being optional in plots diff output

---
 extension/src/cli/dvc/contract.ts    |  8 ++++++--
 extension/src/cli/dvc/reader.test.ts | 25 ++++++++++++++++++++++++-
 extension/src/cli/dvc/reader.ts      | 28 ++++++++++++++++++++++------
 3 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/extension/src/cli/dvc/contract.ts b/extension/src/cli/dvc/contract.ts
index a02b7dc1b4..18d33a3fda 100644
--- a/extension/src/cli/dvc/contract.ts
+++ b/extension/src/cli/dvc/contract.ts
@@ -106,9 +106,13 @@ export type PlotError = {
   source?: string
 } & ErrorContents
 
-export interface PlotsOutput {
-  data: { [path: string]: Plot[] }
+export type RawPlotsOutput = {
+  data?: { [path: string]: Plot[] }
   errors?: PlotError[]
 }
 
+export type PlotsOutput = RawPlotsOutput & {
+  data: { [path: string]: Plot[] }
+}
+
 export type PlotsOutputOrError = PlotsOutput | DvcError
diff --git a/extension/src/cli/dvc/reader.test.ts b/extension/src/cli/dvc/reader.test.ts
index d8bec43897..19c177a0d0 100644
--- a/extension/src/cli/dvc/reader.test.ts
+++ b/extension/src/cli/dvc/reader.test.ts
@@ -180,7 +180,30 @@ describe('CliReader', () => {
       mockedCreateProcess.mockReturnValueOnce(getMockedProcess(''))
 
       const plots = await dvcReader.plotsDiff(cwd, 'HEAD')
-      expect(plots).toStrictEqual({})
+      expect(plots).toStrictEqual({ data: {} })
+    })
+
+    it('should remove an empty errors array from the output', async () => {
+      const cwd = __dirname
+
+      mockedCreateProcess.mockReturnValueOnce(
+        getMockedProcess(JSON.stringify({ errors: [] }))
+      )
+
+      const plots = await dvcReader.plotsDiff(cwd, 'main')
+      expect(plots).toStrictEqual({ data: {} })
+    })
+
+    it('should not remove an errors array with entries from the output', async () => {
+      const cwd = __dirname
+      const errors = [{ msg: 'something went wrong' }]
+
+      mockedCreateProcess.mockReturnValueOnce(
+        getMockedProcess(JSON.stringify({ errors }))
+      )
+
+      const plots = await dvcReader.plotsDiff(cwd, 'main')
+      expect(plots).toStrictEqual({ data: {}, errors })
     })
 
     it('should match the expected output', async () => {
diff --git a/extension/src/cli/dvc/reader.ts b/extension/src/cli/dvc/reader.ts
index 54c3cde02a..e2b18091f8 100644
--- a/extension/src/cli/dvc/reader.ts
+++ b/extension/src/cli/dvc/reader.ts
@@ -14,20 +14,22 @@ import {
   ExperimentsOutput,
   EXPERIMENT_WORKSPACE_ID,
   PlotsOutput,
-  PlotsOutputOrError
+  PlotsOutputOrError,
+  RawPlotsOutput
 } from './contract'
 import { getOptions } from './options'
 import { typeCheckCommands } from '..'
 import { MaybeConsoleError } from '../error'
 import { Logger } from '../../common/logger'
 import { parseNonStandardJson } from '../../util/json'
+import { definedAndNonEmpty } from '../../util/array'
 
 const defaultExperimentsOutput: ExperimentsOutput = {
   [EXPERIMENT_WORKSPACE_ID]: { baseline: {} }
 }
 
 export const isDvcError = <
-  T extends ExperimentsOutput | DataStatusOutput | PlotsOutput
+  T extends ExperimentsOutput | DataStatusOutput | RawPlotsOutput
 >(
   dataOrError: T | DvcError
 ): dataOrError is DvcError =>
@@ -83,11 +85,11 @@ export class DvcReader extends DvcCli {
     return output
   }
 
-  public plotsDiff(
+  public async plotsDiff(
     cwd: string,
     ...revisions: string[]
   ): Promise<PlotsOutputOrError> {
-    return this.readProcessJson<PlotsOutput>(
+    const output = await this.readProcessJson<RawPlotsOutput>(
       cwd,
       Command.PLOTS,
       SubCommand.DIFF,
@@ -96,6 +98,20 @@ export class DvcReader extends DvcCli {
       TEMP_PLOTS_DIR,
       Flag.SPLIT
     )
+    if (isDvcError(output)) {
+      return output
+    }
+
+    const { data, errors } = output
+    const expectedOutput: PlotsOutput = {
+      data: data || {}
+    }
+
+    if (definedAndNonEmpty(errors)) {
+      expectedOutput.errors = errors
+    }
+
+    return expectedOutput
   }
 
   public async root(cwd: string): Promise<string | undefined> {
@@ -125,7 +141,7 @@ export class DvcReader extends DvcCli {
     } catch {}
   }
 
-  private readProcessJson<T extends DataStatusOutput | PlotsOutput>(
+  private readProcessJson<T extends DataStatusOutput | RawPlotsOutput>(
     cwd: string,
     command: Command,
     ...args: Args
@@ -134,7 +150,7 @@ export class DvcReader extends DvcCli {
   }
 
   private async readProcess<
-    T extends DataStatusOutput | ExperimentsOutput | PlotsOutput
+    T extends DataStatusOutput | ExperimentsOutput | RawPlotsOutput
   >(cwd: string, defaultValue: string, ...args: Args): Promise<T | DvcError> {
     try {
       const output =