@@ -21,7 +21,9 @@ import {
2121 writeTsv ,
2222 isPathInProject ,
2323 getPidFromFile ,
24- getEntryFromJsonFile
24+ getEntryFromJsonFile ,
25+ addPlotToDvcYamlFile ,
26+ loadDataFile
2527} from '.'
2628import { dvcDemoPath } from '../test/util'
2729import { DOT_DVC } from '../cli/dvc/constants'
@@ -61,6 +63,83 @@ beforeEach(() => {
6163 jest . resetAllMocks ( )
6264} )
6365
66+ describe ( 'loadDataFile' , ( ) => {
67+ it ( 'should load in csv file contents' , async ( ) => {
68+ const mockCsvContent = [ 'epoch,acc' , '10,0.69' , '11,0.345' ] . join ( '\n' )
69+
70+ mockedReadFileSync . mockReturnValueOnce ( mockCsvContent )
71+
72+ const result = await loadDataFile ( 'values.csv' )
73+
74+ expect ( result ) . toStrictEqual ( [
75+ { acc : 0.69 , epoch : 10 } ,
76+ { acc : 0.345 , epoch : 11 }
77+ ] )
78+ } )
79+
80+ it ( 'should load in json file contents' , async ( ) => {
81+ const mockJsonContent = JSON . stringify ( [
82+ { acc : 0.69 , epoch : 10 } ,
83+ { acc : 0.345 , epoch : 11 }
84+ ] )
85+
86+ mockedReadFileSync . mockReturnValueOnce ( mockJsonContent )
87+
88+ const result = await loadDataFile ( 'values.json' )
89+
90+ expect ( result ) . toStrictEqual ( [
91+ { acc : 0.69 , epoch : 10 } ,
92+ { acc : 0.345 , epoch : 11 }
93+ ] )
94+ } )
95+
96+ it ( 'should load in tsv file contents' , async ( ) => {
97+ const mockTsvContent = [ 'epoch\tacc' , '10\t0.69' , '11\t0.345' ] . join ( '\n' )
98+
99+ mockedReadFileSync . mockReturnValueOnce ( mockTsvContent )
100+
101+ const result = await loadDataFile ( 'values.tsv' )
102+
103+ expect ( result ) . toStrictEqual ( [
104+ { acc : 0.69 , epoch : 10 } ,
105+ { acc : 0.345 , epoch : 11 }
106+ ] )
107+ } )
108+
109+ it ( 'should load in yaml file contents' , async ( ) => {
110+ const mockYamlContent = [
111+ 'stages:' ,
112+ ' train:' ,
113+ ' cmd: python train.py'
114+ ] . join ( '\n' )
115+
116+ mockedReadFileSync . mockReturnValueOnce ( mockYamlContent )
117+
118+ const result = await loadDataFile ( 'dvc.yaml' )
119+
120+ expect ( result ) . toStrictEqual ( {
121+ stages : {
122+ train : {
123+ cmd : 'python train.py'
124+ }
125+ }
126+ } )
127+ } )
128+
129+ it ( 'should catch any errors thrown during file parsing' , async ( ) => {
130+ const dataFiles = [ 'values.csv' , 'file.json' , 'file.tsv' , 'dvc.yaml' ]
131+ mockedReadFileSync . mockImplementation ( ( ) => {
132+ throw new Error ( 'fake error' )
133+ } )
134+
135+ for ( const file of dataFiles ) {
136+ const resultWithErr = await loadDataFile ( file )
137+
138+ expect ( resultWithErr ) . toStrictEqual ( undefined )
139+ }
140+ } )
141+ } )
142+
64143describe ( 'writeJson' , ( ) => {
65144 it ( 'should write unformatted json in given file' , ( ) => {
66145 writeJson ( 'file-name.json' , { array : [ 1 , 2 , 3 ] , number : 1 } )
@@ -436,6 +515,124 @@ describe('findOrCreateDvcYamlFile', () => {
436515 } )
437516} )
438517
518+ describe ( 'addPlotToDvcYamlFile' , ( ) => {
519+ const mockStagesLines = [ 'stages:' , ' train:' , ' cmd: python train.py' ]
520+ const mockPlotsListLines = [
521+ 'plots:' ,
522+ ' - eval/importance.png' ,
523+ ' - Precision-Recall:' ,
524+ ' x: recall' ,
525+ ' y:' ,
526+ ' eval/prc/train.json: precision' ,
527+ ' eval/prc/test.json: precision'
528+ ]
529+ const mockNewPlotLines = [
530+ ' - data.json:' ,
531+ ' template: simple' ,
532+ ' x: epochs' ,
533+ ' y: accuracy'
534+ ]
535+ it ( 'should add a plots list with the new plot if the dvc.yaml file has no plots' , ( ) => {
536+ const mockDvcYamlContent = mockStagesLines . join ( '\n' )
537+ const mockPlotYamlContent = [ '' , 'plots:' , ...mockNewPlotLines , '' ] . join (
538+ '\n'
539+ )
540+ mockedReadFileSync . mockReturnValueOnce ( mockDvcYamlContent )
541+ mockedReadFileSync . mockReturnValueOnce ( mockDvcYamlContent )
542+
543+ addPlotToDvcYamlFile ( '/' , {
544+ dataFile : '/data.json' ,
545+ template : 'simple' ,
546+ x : 'epochs' ,
547+ y : 'accuracy'
548+ } )
549+
550+ expect ( mockedWriteFileSync ) . toHaveBeenCalledWith (
551+ '//dvc.yaml' ,
552+ mockDvcYamlContent + mockPlotYamlContent
553+ )
554+ } )
555+
556+ it ( 'should add the new plot if the dvc.yaml file already has plots' , ( ) => {
557+ const mockDvcYamlContent = [ ...mockPlotsListLines , ...mockStagesLines ]
558+ const mockPlotYamlContent = [ ...mockNewPlotLines , '' ]
559+ mockedReadFileSync . mockReturnValueOnce ( mockDvcYamlContent . join ( '\n' ) )
560+ mockedReadFileSync . mockReturnValueOnce ( mockDvcYamlContent . join ( '\n' ) )
561+
562+ addPlotToDvcYamlFile ( '/' , {
563+ dataFile : '/data.json' ,
564+ template : 'simple' ,
565+ x : 'epochs' ,
566+ y : 'accuracy'
567+ } )
568+
569+ mockDvcYamlContent . splice ( 7 , 0 , ...mockPlotYamlContent )
570+
571+ expect ( mockedWriteFileSync ) . toHaveBeenCalledWith (
572+ '//dvc.yaml' ,
573+ mockDvcYamlContent . join ( '\n' )
574+ )
575+ } )
576+
577+ it ( 'should add a new plot if the dvc.yaml plots list is at bottom of file' , ( ) => {
578+ const mockDvcYamlContent = [ ...mockStagesLines , ...mockPlotsListLines ] . join (
579+ '\n'
580+ )
581+ const mockPlotYamlContent = [ '' , ...mockNewPlotLines , '' ] . join ( '\n' )
582+ mockedReadFileSync . mockReturnValueOnce ( mockDvcYamlContent )
583+ mockedReadFileSync . mockReturnValueOnce ( mockDvcYamlContent )
584+
585+ addPlotToDvcYamlFile ( '/' , {
586+ dataFile : '/data.json' ,
587+ template : 'simple' ,
588+ x : 'epochs' ,
589+ y : 'accuracy'
590+ } )
591+
592+ expect ( mockedWriteFileSync ) . toHaveBeenCalledWith (
593+ '//dvc.yaml' ,
594+ mockDvcYamlContent + mockPlotYamlContent
595+ )
596+ } )
597+
598+ it ( 'should add a new plot with an indent level that matches the dvc.yaml file' , ( ) => {
599+ const mockDvcYamlContent = [
600+ 'stages:' ,
601+ ' train:' ,
602+ ' cmd: python train.py' ,
603+ 'plots:' ,
604+ ' - eval/importance.png' ,
605+ ' - Precision-Recall:' ,
606+ ' x: recall' ,
607+ ' y:' ,
608+ ' eval/prc/train.json: precision' ,
609+ ' eval/prc/test.json: precision'
610+ ] . join ( '\n' )
611+ const mockPlotYamlContent = [
612+ '' ,
613+ ' - data.json:' ,
614+ ' template: simple' ,
615+ ' x: epochs' ,
616+ ' y: accuracy' ,
617+ ''
618+ ] . join ( '\n' )
619+ mockedReadFileSync . mockReturnValueOnce ( mockDvcYamlContent )
620+ mockedReadFileSync . mockReturnValueOnce ( mockDvcYamlContent )
621+
622+ addPlotToDvcYamlFile ( '/' , {
623+ dataFile : '/data.json' ,
624+ template : 'simple' ,
625+ x : 'epochs' ,
626+ y : 'accuracy'
627+ } )
628+
629+ expect ( mockedWriteFileSync ) . toHaveBeenCalledWith (
630+ '//dvc.yaml' ,
631+ mockDvcYamlContent + mockPlotYamlContent
632+ )
633+ } )
634+ } )
635+
439636describe ( 'isPathInProject' , ( ) => {
440637 it ( 'should return true if the path is in the project' , ( ) => {
441638 const path = join ( dvcDemoPath , 'dvc.yaml' )
0 commit comments