Skip to content

Commit 3d9fde3

Browse files
smithjilkssmithjilks
authored
NOISSUE - Enhance CLI (#250)
* Enhance CLI progressbar Signed-off-by: Jilks Smith <smithjilks@gmail.com> * Update cli error and success messages colors Signed-off-by: Jilks Smith <smithjilks@gmail.com> * Update cli emojis Signed-off-by: Jilks Smith <smithjilks@gmail.com> * Add logs for cli interrupt by user Signed-off-by: Jilks Smith <smithjilks@gmail.com> * Remove extra whitespaces Signed-off-by: Jilks Smith <smithjilks@gmail.com> * Update upload data emoji Signed-off-by: Jilks Smith <smithjilks@gmail.com> * Update cli main.go Signed-off-by: Jilks Smith <smithjilks@gmail.com> * Update cli errors Signed-off-by: Jilks Smith <smithjilks@gmail.com> * Update cli Signed-off-by: Jilks Smith <smithjilks@gmail.com> * Update cli Signed-off-by: Jilks Smith <smithjilks@gmail.com> * Update go sum Signed-off-by: Jilks Smith <smithjilks@gmail.com> * Add progressbar tests Signed-off-by: Jilks Smith <smithjilks@gmail.com> * Fix cli cmd error formating Signed-off-by: Jilks Smith <smithjilks@gmail.com> * Add cli datasets, algo and result tests Signed-off-by: Jilks Smith <smithjilks@gmail.com> --------- Signed-off-by: Jilks Smith <smithjilks@gmail.com>
1 parent 63994d7 commit 3d9fde3

File tree

13 files changed

+762
-62
lines changed

13 files changed

+762
-62
lines changed

‎cli/algorithm_test.go‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
// Copyright (c) Ultraviolet
2+
// SPDX-License-Identifier: Apache-2.0
3+
package cli
4+
5+
import (
6+
"bytes"
7+
"crypto/rand"
8+
"crypto/rsa"
9+
"crypto/x509"
10+
"encoding/pem"
11+
"errors"
12+
"log"
13+
"os"
14+
"testing"
15+
16+
"github.com/stretchr/testify/mock"
17+
"github.com/stretchr/testify/require"
18+
"github.com/ultravioletrs/cocos/pkg/sdk/mocks"
19+
)
20+
21+
const algorithmFile = "test_algo_file.py"
22+
23+
func captureLogOutput(f func()) string {
24+
var buf bytes.Buffer
25+
log.SetOutput(&buf)
26+
defer log.SetOutput(os.Stderr)
27+
f()
28+
return buf.String()
29+
}
30+
31+
func generateRSAPrivateKeyFile(fileName string) error {
32+
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
33+
if err != nil {
34+
return err
35+
}
36+
37+
privKeyFile, err := os.Create(fileName)
38+
if err != nil {
39+
return err
40+
}
41+
defer privKeyFile.Close()
42+
43+
privateKeyPEM := &pem.Block{
44+
Type: rsaKeyType,
45+
Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
46+
}
47+
48+
err = pem.Encode(privKeyFile, privateKeyPEM)
49+
if err != nil {
50+
return err
51+
}
52+
53+
return nil
54+
}
55+
56+
func TestAlgorithmCmd_Success(t *testing.T) {
57+
mockSDK := new(mocks.SDK)
58+
mockSDK.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(nil)
59+
testCLI := New(mockSDK)
60+
61+
err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644)
62+
require.NoError(t, err)
63+
64+
err = generateRSAPrivateKeyFile(privateKeyFile)
65+
require.NoError(t, err)
66+
67+
cmd := testCLI.NewAlgorithmCmd()
68+
output := captureLogOutput(func() {
69+
cmd.SetArgs([]string{algorithmFile, privateKeyFile})
70+
err = cmd.Execute()
71+
require.NoError(t, err)
72+
})
73+
74+
require.Contains(t, output, "Successfully uploaded algorithm")
75+
t.Cleanup(func() {
76+
os.Remove(privateKeyFile)
77+
os.Remove(algorithmFile)
78+
})
79+
}
80+
81+
func TestAlgorithmCmd_MissingAlgorithmFile(t *testing.T) {
82+
mockSDK := new(mocks.SDK)
83+
mockSDK.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(nil)
84+
testCLI := New(mockSDK)
85+
86+
cmd := testCLI.NewAlgorithmCmd()
87+
88+
output := captureLogOutput(func() {
89+
cmd.SetArgs([]string{"non_existent_algo_file.py", privateKeyFile})
90+
err := cmd.Execute()
91+
require.NoError(t, err)
92+
})
93+
94+
require.Contains(t, output, "Error reading algorithm file")
95+
}
96+
97+
func TestAlgorithmCmd_MissingPrivateKeyFile(t *testing.T) {
98+
mockSDK := new(mocks.SDK)
99+
mockSDK.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(nil)
100+
testCLI := New(mockSDK)
101+
102+
err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644)
103+
require.NoError(t, err)
104+
105+
cmd := testCLI.NewAlgorithmCmd()
106+
107+
output := captureLogOutput(func() {
108+
cmd.SetArgs([]string{algorithmFile, "non_existent_private_key.pem"})
109+
err = cmd.Execute()
110+
require.NoError(t, err)
111+
})
112+
113+
require.Contains(t, output, "Error reading private key file")
114+
t.Cleanup(func() {
115+
os.Remove(algorithmFile)
116+
})
117+
}
118+
119+
func TestAlgorithmCmd_UploadFailure(t *testing.T) {
120+
mockSDK := new(mocks.SDK)
121+
mockSDK.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(errors.New("failed to upload algorithm due to error"))
122+
testCLI := New(mockSDK)
123+
124+
err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644)
125+
require.NoError(t, err)
126+
127+
err = generateRSAPrivateKeyFile(privateKeyFile)
128+
require.NoError(t, err)
129+
130+
cmd := testCLI.NewAlgorithmCmd()
131+
132+
output := captureLogOutput(func() {
133+
cmd.SetArgs([]string{algorithmFile, privateKeyFile})
134+
err = cmd.Execute()
135+
require.NoError(t, err)
136+
})
137+
138+
require.Contains(t, output, "Failed to upload algorithm")
139+
140+
t.Cleanup(func() {
141+
os.Remove(privateKeyFile)
142+
os.Remove(algorithmFile)
143+
})
144+
}

‎cli/algorithms.go‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"log"
99
"os"
1010

11+
"github.com/fatih/color"
1112
"github.com/spf13/cobra"
1213
"github.com/ultravioletrs/cocos/agent"
1314
"github.com/ultravioletrs/cocos/agent/algorithm"
@@ -35,14 +36,18 @@ func (cli *CLI) NewAlgorithmCmd() *cobra.Command {
3536

3637
algorithm, err := os.ReadFile(algorithmFile)
3738
if err != nil {
38-
log.Fatalf("Error reading algorithm file: %v", err)
39+
msg := color.New(color.FgRed).Sprintf("Error reading algorithm file: %v ❌ ", err)
40+
log.Println(msg)
41+
return
3942
}
4043

4144
var req []byte
4245
if requirementsFile != "" {
4346
req, err = os.ReadFile(requirementsFile)
4447
if err != nil {
45-
log.Fatalf("Error reading requirments file: %v", err)
48+
msg := color.New(color.FgRed).Sprintf("Error reading requirments file: %v ❌ ", err)
49+
log.Println(msg)
50+
return
4651
}
4752
}
4853

@@ -53,7 +58,9 @@ func (cli *CLI) NewAlgorithmCmd() *cobra.Command {
5358

5459
privKeyFile, err := os.ReadFile(args[1])
5560
if err != nil {
56-
log.Fatalf("Error reading private key file: %v", err)
61+
msg := color.New(color.FgRed).Sprintf("Error reading private key file: %v ❌ ", err.Error())
62+
log.Println(msg)
63+
return
5764
}
5865

5966
pemBlock, _ := pem.Decode(privKeyFile)
@@ -63,10 +70,12 @@ func (cli *CLI) NewAlgorithmCmd() *cobra.Command {
6370
ctx := metadata.NewOutgoingContext(cmd.Context(), metadata.New(make(map[string]string)))
6471

6572
if err := cli.agentSDK.Algo(addAlgoMetadata(ctx), algoReq, privKey); err != nil {
66-
log.Fatalf("Error uploading algorithm with error: %v", err)
73+
msg := color.New(color.FgRed).Sprintf("Failed to upload algorithm due to error: %v ❌ ", err.Error())
74+
log.Println(msg)
75+
return
6776
}
6877

69-
log.Println("Successfully uploaded algorithm")
78+
log.Println(color.New(color.FgGreen).Sprint("Successfully uploaded algorithm! ✔ "))
7079
},
7180
}
7281

‎cli/datasets.go‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"os"
1111
"path"
1212

13+
"github.com/fatih/color"
1314
"github.com/spf13/cobra"
1415
"github.com/ultravioletrs/cocos/agent"
1516
"github.com/ultravioletrs/cocos/internal"
@@ -31,20 +32,26 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {
3132

3233
f, err := os.Stat(datasetPath)
3334
if err != nil {
34-
log.Fatalf("Error reading dataset file: %v", err)
35+
msg := color.New(color.FgRed).Sprintf("Error reading dataset file: %v ❌ ", err)
36+
log.Println(msg)
37+
return
3538
}
3639

3740
var dataset []byte
3841

3942
if f.IsDir() {
4043
dataset, err = internal.ZipDirectoryToMemory(datasetPath)
4144
if err != nil {
42-
log.Fatalf("Error zipping dataset directory: %v", err)
45+
msg := color.New(color.FgRed).Sprintf("Error zipping dataset directory: %v ❌ ", err)
46+
log.Println(msg)
47+
return
4348
}
4449
} else {
4550
dataset, err = os.ReadFile(datasetPath)
4651
if err != nil {
47-
log.Fatalf("Error reading dataset file: %v", err)
52+
msg := color.New(color.FgRed).Sprintf("Error reading dataset file: %v ❌ ", err)
53+
log.Println(msg)
54+
return
4855
}
4956
}
5057

@@ -55,7 +62,9 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {
5562

5663
privKeyFile, err := os.ReadFile(args[1])
5764
if err != nil {
58-
log.Fatalf("Error reading private key file: %v", err)
65+
msg := color.New(color.FgRed).Sprintf("Error reading private key file: %v ❌ ", err)
66+
log.Println(msg)
67+
return
5968
}
6069

6170
pemBlock, _ := pem.Decode(privKeyFile)
@@ -64,10 +73,12 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {
6473

6574
ctx := metadata.NewOutgoingContext(cmd.Context(), metadata.New(make(map[string]string)))
6675
if err := cli.agentSDK.Data(addDatasetMetadata(ctx), dataReq, privKey); err != nil {
67-
log.Fatalf("Error uploading dataset: %v", err)
76+
msg := color.New(color.FgRed).Sprintf("Failed to upload dataset due to error: %v ❌ ", err.Error())
77+
log.Println(msg)
78+
return
6879
}
6980

70-
log.Println("Successfully uploaded dataset")
81+
log.Println(color.New(color.FgGreen).Sprint("Successfully uploaded dataset! ✔ "))
7182
},
7283
}
7384

‎cli/datasets_test.go‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
// Copyright (c) Ultraviolet
2+
// SPDX-License-Identifier: Apache-2.0
3+
package cli
4+
5+
import (
6+
"errors"
7+
"os"
8+
"testing"
9+
10+
"github.com/stretchr/testify/mock"
11+
"github.com/stretchr/testify/require"
12+
"github.com/ultravioletrs/cocos/pkg/sdk/mocks"
13+
)
14+
15+
func createTempDatasetFile(content string) (string, error) {
16+
tmpFile, err := os.CreateTemp("", "dataset-*.txt")
17+
if err != nil {
18+
return "", err
19+
}
20+
defer tmpFile.Close()
21+
22+
_, err = tmpFile.WriteString(content)
23+
if err != nil {
24+
return "", err
25+
}
26+
return tmpFile.Name(), nil
27+
}
28+
29+
func TestDatasetsCmd_Success(t *testing.T) {
30+
mockSDK := new(mocks.SDK)
31+
mockSDK.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(nil)
32+
testCLI := New(mockSDK)
33+
34+
datasetFile, err := createTempDatasetFile("test dataset content")
35+
require.NoError(t, err)
36+
37+
err = generateRSAPrivateKeyFile(privateKeyFile)
38+
require.NoError(t, err)
39+
40+
cmd := testCLI.NewDatasetsCmd()
41+
42+
output := captureLogOutput(func() {
43+
cmd.SetArgs([]string{datasetFile, privateKeyFile})
44+
err = cmd.Execute()
45+
require.NoError(t, err)
46+
})
47+
48+
require.Contains(t, output, "Successfully uploaded dataset")
49+
mockSDK.AssertCalled(t, "Data", mock.Anything, mock.Anything, mock.Anything)
50+
51+
t.Cleanup(func() {
52+
os.Remove(datasetFile)
53+
os.Remove(privateKeyFile)
54+
})
55+
}
56+
57+
func TestDatasetsCmd_MissingDatasetFile(t *testing.T) {
58+
mockSDK := new(mocks.SDK)
59+
mockSDK.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(nil)
60+
testCLI := New(mockSDK)
61+
62+
cmd := testCLI.NewDatasetsCmd()
63+
64+
output := captureLogOutput(func() {
65+
cmd.SetArgs([]string{"non_existent_dataset.txt", privateKeyFile})
66+
err := cmd.Execute()
67+
require.NoError(t, err)
68+
})
69+
70+
require.Contains(t, output, "Error reading dataset file")
71+
}
72+
73+
func TestDatasetsCmd_MissingPrivateKeyFile(t *testing.T) {
74+
mockSDK := new(mocks.SDK)
75+
mockSDK.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(nil)
76+
testCLI := New(mockSDK)
77+
78+
datasetFile, err := createTempDatasetFile("test dataset content")
79+
require.NoError(t, err)
80+
81+
cmd := testCLI.NewDatasetsCmd()
82+
83+
output := captureLogOutput(func() {
84+
cmd.SetArgs([]string{datasetFile, "non_existent_private_key.pem"})
85+
err = cmd.Execute()
86+
require.NoError(t, err)
87+
})
88+
89+
require.Contains(t, output, "Error reading private key file")
90+
t.Cleanup(func() {
91+
os.Remove(datasetFile)
92+
})
93+
}
94+
95+
func TestDatasetsCmd_UploadFailure(t *testing.T) {
96+
mockSDK := new(mocks.SDK)
97+
mockSDK.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(errors.New("failed to upload algorithm due to error"))
98+
testCLI := New(mockSDK)
99+
100+
datasetFile, err := createTempDatasetFile("test dataset content")
101+
require.NoError(t, err)
102+
103+
err = generateRSAPrivateKeyFile(privateKeyFile)
104+
require.NoError(t, err)
105+
106+
cmd := testCLI.NewDatasetsCmd()
107+
108+
output := captureLogOutput(func() {
109+
cmd.SetArgs([]string{datasetFile, privateKeyFile})
110+
err = cmd.Execute()
111+
require.NoError(t, err)
112+
})
113+
114+
require.Contains(t, output, "Failed to upload dataset due to error")
115+
t.Cleanup(func() {
116+
os.Remove(datasetFile)
117+
os.Remove(privateKeyFile)
118+
})
119+
}

0 commit comments

Comments
 (0)