mirror of
https://github.com/ultravioletrs/cocos.git
synced 2026-06-23 04:10:25 +00:00
NOISSUE - Improve SDK initialization (#302)
* do not require agent connection for all commands Signed-off-by: Sammy Oina <sammyoina@gmail.com> * improve tests Signed-off-by: Sammy Oina <sammyoina@gmail.com> * more test cases Signed-off-by: Sammy Oina <sammyoina@gmail.com> * blank line Signed-off-by: Sammy Oina <sammyoina@gmail.com> * ignore cmd Signed-off-by: Sammy Oina <sammyoina@gmail.com> --------- Signed-off-by: Sammy Oina <sammyoina@gmail.com>
This commit is contained in:
Sammy Kerata Oina
committed by
GitHub
GitHub
parent
d3636de824
commit
01a619fd2a
+129
-114
@@ -44,122 +44,137 @@ func generateRSAPrivateKeyFile(fileName string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestAlgorithmCmd_Success(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
testCLI := New(mockSDK)
|
||||
func TestAlgorithmCmd(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
setupMock func(*mocks.SDK)
|
||||
setupFiles func() error
|
||||
args []string
|
||||
connectErr error
|
||||
expectedOutput string
|
||||
cleanup func()
|
||||
}{
|
||||
{
|
||||
name: "successful upload",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
},
|
||||
setupFiles: func() error {
|
||||
if err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644); err != nil {
|
||||
return err
|
||||
}
|
||||
return generateRSAPrivateKeyFile(privateKeyFile)
|
||||
},
|
||||
args: []string{algorithmFile, privateKeyFile},
|
||||
expectedOutput: "Successfully uploaded algorithm",
|
||||
cleanup: func() {
|
||||
os.Remove(privateKeyFile)
|
||||
os.Remove(algorithmFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "missing algorithm file",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
},
|
||||
args: []string{"non_existent_algo_file.py", privateKeyFile},
|
||||
expectedOutput: "Error reading algorithm file",
|
||||
},
|
||||
{
|
||||
name: "missing private key file",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
},
|
||||
setupFiles: func() error {
|
||||
return os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644)
|
||||
},
|
||||
args: []string{algorithmFile, "non_existent_private_key.pem"},
|
||||
expectedOutput: "Error reading private key file",
|
||||
cleanup: func() {
|
||||
os.Remove(algorithmFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "upload failure",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(errors.New("failed to upload algorithm due to error"))
|
||||
},
|
||||
setupFiles: func() error {
|
||||
if err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644); err != nil {
|
||||
return err
|
||||
}
|
||||
return generateRSAPrivateKeyFile(privateKeyFile)
|
||||
},
|
||||
args: []string{algorithmFile, privateKeyFile},
|
||||
expectedOutput: "Failed to upload algorithm",
|
||||
cleanup: func() {
|
||||
os.Remove(privateKeyFile)
|
||||
os.Remove(algorithmFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "invalid private key",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
},
|
||||
setupFiles: func() error {
|
||||
if err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644); err != nil {
|
||||
return err
|
||||
}
|
||||
privKeyFile, err := os.Create(privateKeyFile)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer privKeyFile.Close()
|
||||
_, err = privKeyFile.WriteString("invalid private key")
|
||||
return err
|
||||
},
|
||||
args: []string{algorithmFile, privateKeyFile},
|
||||
expectedOutput: "Error decoding private key",
|
||||
cleanup: func() {
|
||||
os.Remove(algorithmFile)
|
||||
os.Remove(privateKeyFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "connection error",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
},
|
||||
connectErr: errors.New("failed to connect to agent"),
|
||||
args: []string{algorithmFile, privateKeyFile},
|
||||
expectedOutput: "Failed to connect to agent",
|
||||
},
|
||||
}
|
||||
|
||||
err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644)
|
||||
require.NoError(t, err)
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
if tt.setupMock != nil {
|
||||
tt.setupMock(mockSDK)
|
||||
}
|
||||
|
||||
err = generateRSAPrivateKeyFile(privateKeyFile)
|
||||
require.NoError(t, err)
|
||||
testCLI := CLI{
|
||||
agentSDK: mockSDK,
|
||||
connectErr: tt.connectErr,
|
||||
}
|
||||
|
||||
cmd := testCLI.NewAlgorithmCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
cmd.SetArgs([]string{algorithmFile, privateKeyFile})
|
||||
err = cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
if tt.setupFiles != nil {
|
||||
err := tt.setupFiles()
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
require.Contains(t, buf.String(), "Successfully uploaded algorithm")
|
||||
t.Cleanup(func() {
|
||||
os.Remove(privateKeyFile)
|
||||
os.Remove(algorithmFile)
|
||||
})
|
||||
}
|
||||
|
||||
func TestAlgorithmCmd_MissingAlgorithmFile(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
testCLI := New(mockSDK)
|
||||
|
||||
cmd := testCLI.NewAlgorithmCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
|
||||
cmd.SetArgs([]string{"non_existent_algo_file.py", privateKeyFile})
|
||||
err := cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), "Error reading algorithm file")
|
||||
}
|
||||
|
||||
func TestAlgorithmCmd_MissingPrivateKeyFile(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
testCLI := New(mockSDK)
|
||||
|
||||
err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644)
|
||||
require.NoError(t, err)
|
||||
|
||||
cmd := testCLI.NewAlgorithmCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
cmd.SetArgs([]string{algorithmFile, "non_existent_private_key.pem"})
|
||||
err = cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), "Error reading private key file")
|
||||
t.Cleanup(func() {
|
||||
os.Remove(algorithmFile)
|
||||
})
|
||||
}
|
||||
|
||||
func TestAlgorithmCmd_UploadFailure(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(errors.New("failed to upload algorithm due to error"))
|
||||
testCLI := New(mockSDK)
|
||||
|
||||
err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644)
|
||||
require.NoError(t, err)
|
||||
|
||||
err = generateRSAPrivateKeyFile(privateKeyFile)
|
||||
require.NoError(t, err)
|
||||
|
||||
cmd := testCLI.NewAlgorithmCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
|
||||
cmd.SetArgs([]string{algorithmFile, privateKeyFile})
|
||||
err = cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), "Failed to upload algorithm")
|
||||
|
||||
t.Cleanup(func() {
|
||||
os.Remove(privateKeyFile)
|
||||
os.Remove(algorithmFile)
|
||||
})
|
||||
}
|
||||
|
||||
func TestAlgorithmCmd_InvalidPrivateKey(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
testCLI := New(mockSDK)
|
||||
|
||||
err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644)
|
||||
require.NoError(t, err)
|
||||
|
||||
privKeyFile, err := os.Create(privateKeyFile)
|
||||
require.NoError(t, err)
|
||||
defer privKeyFile.Close()
|
||||
|
||||
_, err = privKeyFile.WriteString("invalid private key")
|
||||
require.NoError(t, err)
|
||||
|
||||
cmd := testCLI.NewAlgorithmCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
|
||||
cmd.SetArgs([]string{algorithmFile, privateKeyFile})
|
||||
err = cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), "Error decoding private key")
|
||||
|
||||
t.Cleanup(func() {
|
||||
os.Remove(algorithmFile)
|
||||
os.Remove(privateKeyFile)
|
||||
})
|
||||
cmd := testCLI.NewAlgorithmCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
cmd.SetArgs(tt.args)
|
||||
err := cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), tt.expectedOutput)
|
||||
|
||||
if tt.cleanup != nil {
|
||||
tt.cleanup()
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -29,6 +29,11 @@ func (cli *CLI) NewAlgorithmCmd() *cobra.Command {
|
||||
Example: "algo <algo_file> <private_key_file_path>",
|
||||
Args: cobra.ExactArgs(2),
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
if cli.connectErr != nil {
|
||||
printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
|
||||
return
|
||||
}
|
||||
|
||||
algorithmFile := args[0]
|
||||
|
||||
cmd.Println("Uploading algorithm file:", algorithmFile)
|
||||
|
||||
+12
-7
@@ -126,9 +126,9 @@ func (cli *CLI) NewAttestationCmd() *cobra.Command {
|
||||
Use: "attestation [command]",
|
||||
Short: "Get and validate attestations",
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
fmt.Printf("Get and validate attestations\n\n")
|
||||
fmt.Printf("Usage:\n %s [command]\n\n", cmd.CommandPath())
|
||||
fmt.Printf("Available Commands:\n")
|
||||
cmd.Printf("Get and validate attestations\n\n")
|
||||
cmd.Printf("Usage:\n %s [command]\n\n", cmd.CommandPath())
|
||||
cmd.Printf("Available Commands:\n")
|
||||
|
||||
// Filter out "completion" command
|
||||
availableCommands := make([]*cobra.Command, 0)
|
||||
@@ -139,14 +139,14 @@ func (cli *CLI) NewAttestationCmd() *cobra.Command {
|
||||
}
|
||||
|
||||
for _, subCmd := range availableCommands {
|
||||
fmt.Printf(" %-15s%s\n", subCmd.Name(), subCmd.Short)
|
||||
cmd.Printf(" %-15s%s\n", subCmd.Name(), subCmd.Short)
|
||||
}
|
||||
|
||||
fmt.Printf("\nFlags:\n")
|
||||
cmd.Printf("\nFlags:\n")
|
||||
cmd.Flags().VisitAll(func(flag *pflag.Flag) {
|
||||
fmt.Printf(" -%s, --%s %s\n", flag.Shorthand, flag.Name, flag.Usage)
|
||||
cmd.Printf(" -%s, --%s %s\n", flag.Shorthand, flag.Name, flag.Usage)
|
||||
})
|
||||
fmt.Printf("\nUse \"%s [command] --help\" for more information about a command.\n", cmd.CommandPath())
|
||||
cmd.Printf("\nUse \"%s [command] --help\" for more information about a command.\n", cmd.CommandPath())
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -158,6 +158,11 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
|
||||
Example: "get <report_data>",
|
||||
Args: cobra.ExactArgs(1),
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
if cli.connectErr != nil {
|
||||
printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
|
||||
return
|
||||
}
|
||||
|
||||
cmd.Println("Getting attestation")
|
||||
|
||||
reportData, err := hex.DecodeString(args[0])
|
||||
|
||||
+99
-17
@@ -27,29 +27,111 @@ func TestNewAttestationCmd(t *testing.T) {
|
||||
|
||||
assert.Equal(t, "attestation [command]", cmd.Use)
|
||||
assert.Equal(t, "Get and validate attestations", cmd.Short)
|
||||
}
|
||||
|
||||
func TestNewGetAttestationCmd(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
cli := &CLI{agentSDK: mockSDK}
|
||||
cmd := cli.NewGetAttestationCmd()
|
||||
var buf bytes.Buffer
|
||||
|
||||
cmd.SetOutput(&buf)
|
||||
|
||||
assert.Equal(t, "get", cmd.Use)
|
||||
assert.Equal(t, "Retrieve attestation information from agent. Report data expected in hex enoded string of length 64 bytes.", cmd.Short)
|
||||
|
||||
reportData := bytes.Repeat([]byte{0x01}, agent.ReportDataSize)
|
||||
mockSDK.On("Attestation", mock.Anything, [agent.ReportDataSize]byte(reportData)).Return([]byte("mock attestation"), nil)
|
||||
|
||||
cmd.SetArgs([]string{hex.EncodeToString(reportData)})
|
||||
cmd.SetOut(&buf)
|
||||
err := cmd.Execute()
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, buf.String(), "Get and validate attestations")
|
||||
}
|
||||
|
||||
assert.Contains(t, buf.String(), "Attestation result retrieved and saved successfully!")
|
||||
func TestNewGetAttestationCmdN(t *testing.T) {
|
||||
validattestation, err := os.ReadFile("../attestation.bin")
|
||||
require.NoError(t, err)
|
||||
testCases := []struct {
|
||||
name string
|
||||
args []string
|
||||
mockResponse []byte
|
||||
mockError error
|
||||
expectedErr string
|
||||
expectedOut string
|
||||
}{
|
||||
{
|
||||
name: "successful attestation retrieval",
|
||||
args: []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize))},
|
||||
mockResponse: []byte("mock attestation"),
|
||||
mockError: nil,
|
||||
expectedOut: "Attestation result retrieved and saved successfully!",
|
||||
},
|
||||
{
|
||||
name: "invalid report data (decoding error)",
|
||||
args: []string{"invalid"},
|
||||
mockResponse: nil,
|
||||
mockError: errors.New("error"),
|
||||
expectedErr: "Error decoding report data",
|
||||
},
|
||||
{
|
||||
name: "invalid report data size",
|
||||
args: []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, 32))},
|
||||
mockResponse: nil,
|
||||
mockError: errors.New("error"),
|
||||
expectedErr: "report data must be a hex encoded string of length 64 bytes",
|
||||
},
|
||||
{
|
||||
name: "invalid report data hex",
|
||||
args: []string{"invalid"},
|
||||
mockResponse: nil,
|
||||
mockError: errors.New("error"),
|
||||
expectedErr: "Error decoding report data",
|
||||
},
|
||||
{
|
||||
name: "failed to get attestation",
|
||||
args: []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize))},
|
||||
mockResponse: nil,
|
||||
mockError: errors.New("error"),
|
||||
expectedErr: "Failed to get attestation due to error",
|
||||
},
|
||||
{
|
||||
name: "JSON report error",
|
||||
args: []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize)), "--json"},
|
||||
mockResponse: []byte("mock attestation"),
|
||||
mockError: nil,
|
||||
expectedErr: "Error converting attestation to json",
|
||||
},
|
||||
{
|
||||
name: "successful JSON report",
|
||||
args: []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize)), "--json"},
|
||||
mockResponse: validattestation,
|
||||
mockError: nil,
|
||||
expectedOut: "Attestation result retrieved and saved successfully!",
|
||||
},
|
||||
{
|
||||
name: "connection error",
|
||||
args: []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize))},
|
||||
mockResponse: nil,
|
||||
mockError: errors.New("failed to connect to agent"),
|
||||
expectedErr: "Failed to connect to agent",
|
||||
},
|
||||
}
|
||||
|
||||
os.Remove(attestationFilePath)
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Cleanup(func() {
|
||||
os.Remove(attestationFilePath)
|
||||
os.Remove(attestationJson)
|
||||
})
|
||||
mockSDK := new(mocks.SDK)
|
||||
cli := &CLI{agentSDK: mockSDK}
|
||||
if tc.name == "connection error" {
|
||||
cli.connectErr = errors.New("failed to connect to agent")
|
||||
}
|
||||
cmd := cli.NewGetAttestationCmd()
|
||||
var buf bytes.Buffer
|
||||
cmd.SetOutput(&buf)
|
||||
|
||||
mockSDK.On("Attestation", mock.Anything, [agent.ReportDataSize]byte(bytes.Repeat([]byte{0x01}, agent.ReportDataSize))).Return(tc.mockResponse, tc.mockError)
|
||||
|
||||
cmd.SetArgs(tc.args)
|
||||
err := cmd.Execute()
|
||||
|
||||
if tc.expectedErr != "" {
|
||||
assert.Contains(t, buf.String(), tc.expectedErr)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, buf.String(), tc.expectedOut)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewValidateAttestationValidationCmd(t *testing.T) {
|
||||
|
||||
@@ -26,6 +26,11 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {
|
||||
Example: "data <dataset_path> <private_key_file_path>",
|
||||
Args: cobra.ExactArgs(2),
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
if cli.connectErr != nil {
|
||||
printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
|
||||
return
|
||||
}
|
||||
|
||||
datasetPath := args[0]
|
||||
|
||||
cmd.Println("Uploading dataset:", datasetPath)
|
||||
|
||||
+122
-110
@@ -27,120 +27,132 @@ func createTempDatasetFile(content string) (string, error) {
|
||||
return tmpFile.Name(), nil
|
||||
}
|
||||
|
||||
func TestDatasetsCmd_Success(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
testCLI := New(mockSDK)
|
||||
func TestDatasetsCmd(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
setupMock func(*mocks.SDK)
|
||||
setupFiles func() (string, error)
|
||||
connectErr error
|
||||
expectedOutput string
|
||||
cleanup func(string, string)
|
||||
}{
|
||||
{
|
||||
name: "successful upload",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
},
|
||||
setupFiles: func() (string, error) {
|
||||
datasetFile, err := createTempDatasetFile("test dataset content")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
err = generateRSAPrivateKeyFile(privateKeyFile)
|
||||
return datasetFile, err
|
||||
},
|
||||
expectedOutput: "Successfully uploaded dataset",
|
||||
cleanup: func(datasetFile, privateKeyFile string) {
|
||||
os.Remove(datasetFile)
|
||||
os.Remove(privateKeyFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "missing dataset file",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
},
|
||||
setupFiles: func() (string, error) {
|
||||
return "", nil
|
||||
},
|
||||
expectedOutput: "Error reading dataset file",
|
||||
},
|
||||
{
|
||||
name: "missing private key file",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
},
|
||||
setupFiles: func() (string, error) {
|
||||
return createTempDatasetFile("test dataset content")
|
||||
},
|
||||
expectedOutput: "Error reading private key file",
|
||||
cleanup: func(datasetFile, _ string) {
|
||||
os.Remove(datasetFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "upload failure",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(errors.New("failed to upload algorithm due to error"))
|
||||
},
|
||||
setupFiles: func() (string, error) {
|
||||
datasetFile, err := createTempDatasetFile("test dataset content")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
err = generateRSAPrivateKeyFile(privateKeyFile)
|
||||
return datasetFile, err
|
||||
},
|
||||
expectedOutput: "Failed to upload dataset due to error",
|
||||
cleanup: func(datasetFile, privateKeyFile string) {
|
||||
os.Remove(datasetFile)
|
||||
os.Remove(privateKeyFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "invalid private key",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
},
|
||||
setupFiles: func() (string, error) {
|
||||
datasetFile, err := createTempDatasetFile("test dataset content")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
err = os.WriteFile(privateKeyFile, []byte("invalid private key"), 0o644)
|
||||
return datasetFile, err
|
||||
},
|
||||
expectedOutput: "Error decoding private key",
|
||||
cleanup: func(datasetFile, privateKeyFile string) {
|
||||
os.Remove(datasetFile)
|
||||
os.Remove(privateKeyFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "connection error",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
},
|
||||
setupFiles: func() (string, error) { return "", nil },
|
||||
connectErr: errors.New("failed to connect to agent"),
|
||||
expectedOutput: "Failed to connect to agent",
|
||||
},
|
||||
}
|
||||
|
||||
datasetFile, err := createTempDatasetFile("test dataset content")
|
||||
require.NoError(t, err)
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
if tt.setupMock != nil {
|
||||
tt.setupMock(mockSDK)
|
||||
}
|
||||
|
||||
err = generateRSAPrivateKeyFile(privateKeyFile)
|
||||
require.NoError(t, err)
|
||||
testCLI := CLI{
|
||||
agentSDK: mockSDK,
|
||||
connectErr: tt.connectErr,
|
||||
}
|
||||
|
||||
cmd := testCLI.NewDatasetsCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
datasetFile, err := tt.setupFiles()
|
||||
require.NoError(t, err)
|
||||
|
||||
cmd.SetArgs([]string{datasetFile, privateKeyFile})
|
||||
err = cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
cmd := testCLI.NewDatasetsCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
cmd.SetArgs([]string{datasetFile, privateKeyFile})
|
||||
err = cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), "Successfully uploaded dataset")
|
||||
mockSDK.AssertCalled(t, "Data", mock.Anything, mock.Anything, mock.Anything)
|
||||
require.Contains(t, buf.String(), tt.expectedOutput)
|
||||
|
||||
t.Cleanup(func() {
|
||||
os.Remove(datasetFile)
|
||||
os.Remove(privateKeyFile)
|
||||
})
|
||||
}
|
||||
|
||||
func TestDatasetsCmd_MissingDatasetFile(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
testCLI := New(mockSDK)
|
||||
|
||||
cmd := testCLI.NewDatasetsCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
|
||||
cmd.SetArgs([]string{"non_existent_dataset.txt", privateKeyFile})
|
||||
err := cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), "Error reading dataset file")
|
||||
}
|
||||
|
||||
func TestDatasetsCmd_MissingPrivateKeyFile(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
testCLI := New(mockSDK)
|
||||
|
||||
datasetFile, err := createTempDatasetFile("test dataset content")
|
||||
require.NoError(t, err)
|
||||
|
||||
cmd := testCLI.NewDatasetsCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
|
||||
cmd.SetArgs([]string{datasetFile, "non_existent_private_key.pem"})
|
||||
err = cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), "Error reading private key file")
|
||||
t.Cleanup(func() {
|
||||
os.Remove(datasetFile)
|
||||
})
|
||||
}
|
||||
|
||||
func TestDatasetsCmd_UploadFailure(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(errors.New("failed to upload algorithm due to error"))
|
||||
testCLI := New(mockSDK)
|
||||
|
||||
datasetFile, err := createTempDatasetFile("test dataset content")
|
||||
require.NoError(t, err)
|
||||
|
||||
err = generateRSAPrivateKeyFile(privateKeyFile)
|
||||
require.NoError(t, err)
|
||||
|
||||
cmd := testCLI.NewDatasetsCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
|
||||
cmd.SetArgs([]string{datasetFile, privateKeyFile})
|
||||
err = cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), "Failed to upload dataset due to error")
|
||||
t.Cleanup(func() {
|
||||
os.Remove(datasetFile)
|
||||
os.Remove(privateKeyFile)
|
||||
})
|
||||
}
|
||||
|
||||
func TestDatasetsCmd_InvalidPrivateKey(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(nil)
|
||||
testCLI := New(mockSDK)
|
||||
|
||||
datasetFile, err := createTempDatasetFile("test dataset content")
|
||||
require.NoError(t, err)
|
||||
|
||||
err = os.WriteFile(privateKeyFile, []byte("invalid private key"), 0o644)
|
||||
require.NoError(t, err)
|
||||
|
||||
cmd := testCLI.NewDatasetsCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
|
||||
cmd.SetArgs([]string{datasetFile, privateKeyFile})
|
||||
err = cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), "Error decoding private key")
|
||||
t.Cleanup(func() {
|
||||
os.Remove(datasetFile)
|
||||
os.Remove(privateKeyFile)
|
||||
})
|
||||
if tt.cleanup != nil {
|
||||
tt.cleanup(datasetFile, privateKeyFile)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -23,6 +23,11 @@ func (cli *CLI) NewResultsCmd() *cobra.Command {
|
||||
Example: "result <private_key_file_path>",
|
||||
Args: cobra.ExactArgs(1),
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
if cli.connectErr != nil {
|
||||
printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
|
||||
return
|
||||
}
|
||||
|
||||
cmd.Println("⏳ Retrieving computation result file")
|
||||
|
||||
privKeyFile, err := os.ReadFile(args[0])
|
||||
|
||||
+112
-102
@@ -18,43 +18,10 @@ import (
|
||||
|
||||
const compResult = "Test computation result"
|
||||
|
||||
func TestResultsCmd_Success(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Result", mock.Anything, mock.Anything).Return([]byte(compResult), nil)
|
||||
testCLI := New(mockSDK)
|
||||
|
||||
err := generateRSAPrivateKeyFile(privateKeyFile)
|
||||
require.NoError(t, err)
|
||||
|
||||
cmd := testCLI.NewResultsCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
cmd.SetArgs([]string{privateKeyFile})
|
||||
err = cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), "Computation result retrieved and saved successfully")
|
||||
|
||||
files, err := filepath.Glob("results*.zip")
|
||||
require.NoError(t, err)
|
||||
require.Len(t, files, 1)
|
||||
|
||||
resultFile, err := os.ReadFile(files[0])
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, compResult, string(resultFile))
|
||||
|
||||
t.Cleanup(func() {
|
||||
for _, file := range files {
|
||||
os.Remove(file)
|
||||
}
|
||||
os.Remove(privateKeyFile)
|
||||
})
|
||||
}
|
||||
|
||||
func TestResultsCmd_MultipleExecutions(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Result", mock.Anything, mock.Anything).Return([]byte(compResult), nil)
|
||||
testCLI := New(mockSDK)
|
||||
testCLI := CLI{agentSDK: mockSDK}
|
||||
|
||||
err := generateRSAPrivateKeyFile(privateKeyFile)
|
||||
require.NoError(t, err)
|
||||
@@ -83,77 +50,10 @@ func TestResultsCmd_MultipleExecutions(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestResultsCmd_MissingPrivateKeyFile(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Result", mock.Anything, mock.Anything).Return([]byte(compResult), nil)
|
||||
testCLI := New(mockSDK)
|
||||
|
||||
cmd := testCLI.NewResultsCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
cmd.SetArgs([]string{"non_existent_private_key.pem"})
|
||||
err := cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), "Error reading private key file")
|
||||
}
|
||||
|
||||
func TestResultsCmd_ResultFailure(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Result", mock.Anything, mock.Anything).Return(nil, errors.New("error retrieving computation result"))
|
||||
testCLI := New(mockSDK)
|
||||
|
||||
err := generateRSAPrivateKeyFile(privateKeyFile)
|
||||
require.NoError(t, err)
|
||||
|
||||
cmd := testCLI.NewResultsCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
cmd.SetArgs([]string{privateKeyFile})
|
||||
err = cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), "error retrieving computation result")
|
||||
mockSDK.AssertCalled(t, "Result", mock.Anything, mock.Anything)
|
||||
t.Cleanup(func() {
|
||||
os.Remove(privateKeyFile)
|
||||
})
|
||||
}
|
||||
|
||||
func TestResultsCmd_SaveFailure(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Result", mock.Anything, mock.Anything).Return([]byte(compResult), nil)
|
||||
testCLI := New(mockSDK)
|
||||
|
||||
err := generateRSAPrivateKeyFile(privateKeyFile)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Simulate failure in saving the result file by making all files read-only
|
||||
err = os.Chmod(".", 0o555)
|
||||
require.NoError(t, err)
|
||||
|
||||
cmd := testCLI.NewResultsCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
cmd.SetArgs([]string{privateKeyFile})
|
||||
err = cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), "Error saving computation result file")
|
||||
mockSDK.AssertCalled(t, "Result", mock.Anything, mock.Anything)
|
||||
|
||||
t.Cleanup(func() {
|
||||
err := os.Chmod(".", 0o755)
|
||||
require.NoError(t, err)
|
||||
err = os.Remove(privateKeyFile)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestResultsCmd_InvalidPrivateKey(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
mockSDK.On("Result", mock.Anything, mock.Anything).Return([]byte(compResult), nil)
|
||||
testCLI := New(mockSDK)
|
||||
testCLI := CLI{agentSDK: mockSDK}
|
||||
|
||||
invalidPrivateKey, err := os.CreateTemp("", "invalid_private_key.pem")
|
||||
require.NoError(t, err)
|
||||
@@ -198,3 +98,113 @@ func TestGetUniqueFilePath(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "test_3.txt", path)
|
||||
}
|
||||
|
||||
func TestResultsCmd(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
setupMock func(*mocks.SDK)
|
||||
setupFiles func() (string, error)
|
||||
connectErr error
|
||||
expectedOutput string
|
||||
cleanup func()
|
||||
}{
|
||||
{
|
||||
name: "successful result retrieval",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Result", mock.Anything, mock.Anything).Return([]byte(compResult), nil)
|
||||
},
|
||||
setupFiles: func() (string, error) {
|
||||
return privateKeyFile, generateRSAPrivateKeyFile(privateKeyFile)
|
||||
},
|
||||
expectedOutput: "Computation result retrieved and saved successfully",
|
||||
cleanup: func() {
|
||||
files, _ := filepath.Glob("results*.zip")
|
||||
for _, file := range files {
|
||||
os.Remove(file)
|
||||
}
|
||||
os.Remove(privateKeyFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "missing private key file",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Result", mock.Anything, mock.Anything).Return([]byte(compResult), nil)
|
||||
},
|
||||
setupFiles: func() (string, error) {
|
||||
return "non_existent_private_key.pem", nil
|
||||
},
|
||||
expectedOutput: "Error reading private key file",
|
||||
},
|
||||
{
|
||||
name: "result retrieval failure",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Result", mock.Anything, mock.Anything).Return(nil, errors.New("error retrieving computation result"))
|
||||
},
|
||||
setupFiles: func() (string, error) {
|
||||
return privateKeyFile, generateRSAPrivateKeyFile(privateKeyFile)
|
||||
},
|
||||
expectedOutput: "error retrieving computation result",
|
||||
cleanup: func() {
|
||||
os.Remove(privateKeyFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "save failure",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
m.On("Result", mock.Anything, mock.Anything).Return([]byte(compResult), nil)
|
||||
},
|
||||
setupFiles: func() (string, error) {
|
||||
err := generateRSAPrivateKeyFile(privateKeyFile)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
// Simulate failure in saving the result file by making all files read-only
|
||||
return privateKeyFile, os.Chmod(".", 0o555)
|
||||
},
|
||||
expectedOutput: "Error saving computation result file",
|
||||
cleanup: func() {
|
||||
err := os.Chmod(".", 0o755)
|
||||
require.NoError(t, err)
|
||||
os.Remove(privateKeyFile)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "connection error",
|
||||
setupMock: func(m *mocks.SDK) {
|
||||
},
|
||||
setupFiles: func() (string, error) { return "", nil },
|
||||
connectErr: errors.New("failed to connect to agent"),
|
||||
expectedOutput: "Failed to connect to agent",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
mockSDK := new(mocks.SDK)
|
||||
if tt.setupMock != nil {
|
||||
tt.setupMock(mockSDK)
|
||||
}
|
||||
|
||||
testCLI := CLI{
|
||||
agentSDK: mockSDK,
|
||||
connectErr: tt.connectErr,
|
||||
}
|
||||
|
||||
file, err := tt.setupFiles()
|
||||
require.NoError(t, err)
|
||||
|
||||
cmd := testCLI.NewResultsCmd()
|
||||
buf := new(bytes.Buffer)
|
||||
cmd.SetOut(buf)
|
||||
cmd.SetArgs([]string{file})
|
||||
err = cmd.Execute()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Contains(t, buf.String(), tt.expectedOutput)
|
||||
|
||||
if tt.cleanup != nil {
|
||||
tt.cleanup()
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
+29
-4
@@ -2,16 +2,41 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
package cli
|
||||
|
||||
import "github.com/ultravioletrs/cocos/pkg/sdk"
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/ultravioletrs/cocos/pkg/clients/grpc"
|
||||
"github.com/ultravioletrs/cocos/pkg/clients/grpc/agent"
|
||||
"github.com/ultravioletrs/cocos/pkg/sdk"
|
||||
)
|
||||
|
||||
var Verbose bool
|
||||
|
||||
type CLI struct {
|
||||
agentSDK sdk.SDK
|
||||
agentSDK sdk.SDK
|
||||
config grpc.Config
|
||||
client grpc.Client
|
||||
connectErr error
|
||||
}
|
||||
|
||||
func New(agentSDK sdk.SDK) *CLI {
|
||||
func New(config grpc.Config) *CLI {
|
||||
return &CLI{
|
||||
agentSDK: agentSDK,
|
||||
config: config,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *CLI) InitializeSDK() error {
|
||||
agentGRPCClient, agentClient, err := agent.NewAgentClient(context.Background(), c.config)
|
||||
if err != nil {
|
||||
c.connectErr = err
|
||||
return err
|
||||
}
|
||||
c.client = agentGRPCClient
|
||||
|
||||
c.agentSDK = sdk.NewAgentSDK(agentClient)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *CLI) Close() {
|
||||
c.client.Close()
|
||||
}
|
||||
|
||||
+4
-13
@@ -3,7 +3,6 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/signal"
|
||||
@@ -16,8 +15,6 @@ import (
|
||||
"github.com/spf13/pflag"
|
||||
"github.com/ultravioletrs/cocos/cli"
|
||||
"github.com/ultravioletrs/cocos/pkg/clients/grpc"
|
||||
"github.com/ultravioletrs/cocos/pkg/clients/grpc/agent"
|
||||
"github.com/ultravioletrs/cocos/pkg/sdk"
|
||||
cmd "github.com/virtee/sev-snp-measure-go/sevsnpmeasure/cmd"
|
||||
)
|
||||
|
||||
@@ -101,17 +98,11 @@ func main() {
|
||||
return
|
||||
}
|
||||
|
||||
agentGRPCClient, agentClient, err := agent.NewAgentClient(context.Background(), agentGRPCConfig)
|
||||
if err != nil {
|
||||
message := color.New(color.FgRed).Sprintf("failed to create %s gRPC client : %s", svcName, err)
|
||||
rootCmd.Println(message)
|
||||
return
|
||||
cliSVC := cli.New(agentGRPCConfig)
|
||||
|
||||
if err := cliSVC.InitializeSDK(); err == nil {
|
||||
defer cliSVC.Close()
|
||||
}
|
||||
defer agentGRPCClient.Close()
|
||||
|
||||
agentSDK := sdk.NewAgentSDK(agentClient)
|
||||
|
||||
cliSVC := cli.New(agentSDK)
|
||||
|
||||
rootCmd.PersistentFlags().BoolVarP(&cli.Verbose, "verbose", "v", false, "Enable verbose output")
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
coverage:
|
||||
ignore:
|
||||
- "test/*"
|
||||
- "cmd/*"
|
||||
- "**/mocks/**"
|
||||
- "mocks/**"
|
||||
- "**/*.pb.go"
|
||||
|
||||
@@ -58,7 +58,7 @@ require (
|
||||
github.com/go-logr/stdr v1.2.2 // indirect
|
||||
github.com/golang/protobuf v1.5.4 // indirect
|
||||
github.com/google/go-configfs-tsm v0.2.2 // indirect
|
||||
github.com/google/logger v1.1.1 // indirect
|
||||
github.com/google/logger v1.1.1
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0 // indirect
|
||||
github.com/inconshreveable/mousetrap v1.1.0 // indirect
|
||||
|
||||
Reference in New Issue
Block a user