Skip to content

Commit

Permalink
Refactoring, improve code coverage (#72)
Browse files Browse the repository at this point in the history
* Refactoring, improve code coverage

* Add unit test for gelAllFilesNameByDirName

* Add codecov coverage into README.md

* Improve coverage readFileBytesByFilePath
  • Loading branch information
mariocandela committed Oct 8, 2023
1 parent 1f48f4d commit 07ffdd8
Show file tree
Hide file tree
Showing 9 changed files with 224 additions and 78 deletions.
3 changes: 3 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

[![CI](https://github.com/mariocandela/beelzebub/actions/workflows/ci.yml/badge.svg)](https://github.com/mariocandela/beelzebub/actions/workflows/ci.yml) [![Docker](https://github.com/mariocandela/beelzebub/actions/workflows/docker-image.yml/badge.svg)](https://github.com/mariocandela/beelzebub/actions/workflows/docker-image.yml) [![codeql](https://github.com/mariocandela/beelzebub/actions/workflows/codeql.yml/badge.svg)](https://github.com/mariocandela/beelzebub/actions/workflows/codeql.yml)
[![Go Report Card](https://goreportcard.com/badge/github.com/mariocandela/beelzebub)](https://goreportcard.com/report/github.com/mariocandela/beelzebub)
[![codecov](https://codecov.io/gh/mariocandela/beelzebub/graph/badge.svg?token=8XTK7D4WHE)](https://codecov.io/gh/mariocandela/beelzebub)

## Overview

Expand Down Expand Up @@ -91,7 +92,9 @@ $ make test.unit
To run integration tests:

```bash
$ make test.dependencies.start
$ make test.integration
$ make test.dependencies.down
```

## Key Features
Expand Down
3 changes: 1 addition & 2 deletions parser/configurations_parser.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ package parser

import (
"fmt"
"io/ioutil"
"os"
"path/filepath"
"strings"
Expand Down Expand Up @@ -127,7 +126,7 @@ func (bp configurationsParser) ReadConfigurationsServices() ([]BeelzebubServiceC
}

func gelAllFilesNameByDirName(dirName string) ([]string, error) {
files, err := ioutil.ReadDir(dirName)
files, err := os.ReadDir(dirName)
if err != nil {
return nil, err
}
Expand Down
52 changes: 52 additions & 0 deletions parser/configurations_parser_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package parser

import (
"errors"
"os"
"testing"

"github.com/stretchr/testify/assert"
Expand Down Expand Up @@ -118,3 +119,54 @@ func TestReadConfigurationsServicesValid(t *testing.T) {
assert.Equal(t, len(firstBeelzebubServiceConfiguration.Commands[0].Headers), 1)
assert.Equal(t, firstBeelzebubServiceConfiguration.Commands[0].Headers[0], "Content-Type: text/html")
}

func TestGelAllFilesNameByDirName(t *testing.T) {

var dir = t.TempDir()

files, err := gelAllFilesNameByDirName(dir)

assert.Nil(t, err)
assert.Equal(t, 0, len(files))
}

func TestGelAllFilesNameByDirNameFiles(t *testing.T) {

var dir = t.TempDir()

testFiles := []string{"file1.yaml", "file2.yaml", "file3.txt", "subdir", "file4.yaml"}
for _, filename := range testFiles {
filePath := dir + "/" + filename
file, err := os.Create(filePath)
assert.NoError(t, err)
file.Close()
}

files, err := gelAllFilesNameByDirName(dir)

assert.Nil(t, err)
assert.Equal(t, 3, len(files))
}

func TestGelAllFilesNameByDirNameError(t *testing.T) {

files, err := gelAllFilesNameByDirName("nosuchfile")

assert.Nil(t, files)
assert.Equal(t, "open nosuchfile: no such file or directory", err.Error())
}

func TestReadFileBytesByFilePath(t *testing.T) {

var dir = t.TempDir()
filePath := dir + "/test.yaml"

f, err := os.Create(filePath)
assert.NoError(t, err)
f.Close()

bytes, err := readFileBytesByFilePath(filePath)
assert.NoError(t, err)

assert.Equal(t, "", string(bytes))
}
40 changes: 21 additions & 19 deletions plugins/openai-gpt.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,33 +5,27 @@ import (
"errors"
"fmt"
"strings"

log "github.com/sirupsen/logrus"

"github.com/go-resty/resty/v2"
)

const (
// Reference: https://www.engraved.blog/building-a-virtual-machine-inside/
promptVirtualizeLinuxTerminal = "I want you to act as a Linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. Do no write explanations. Do not type commands unless I instruct you to do so.\n\nA:pwd\n\nQ:/home/user\n\n"
ChatGPTPluginName = "OpenAIGPTLinuxTerminal"
openAIGPTEndpoint = "https://api.openai.com/v1/completions"
)
ChatGPTPluginName = "OpenAIGPTLinuxTerminal"
openAIGPTEndpoint = "https://api.openai.com/v1/completions"
)

type History struct {
Input, Output string
}

type OpenAIGPTVirtualTerminal struct {
Histories []History
OpenAPIChatGPTSecretKey string
client *resty.Client
}

func (openAIGPTVirtualTerminal *OpenAIGPTVirtualTerminal) InjectDependency() {
if openAIGPTVirtualTerminal.client == nil {
openAIGPTVirtualTerminal.client = resty.New()
}
type openAIGPTVirtualTerminal struct {
Histories []History
openAIKey string
client *resty.Client
}

type Choice struct {
Expand Down Expand Up @@ -65,6 +59,14 @@ type gptRequest struct {
Stop []string `json:"stop"`
}

func Init(history []History, openAIKey string) *openAIGPTVirtualTerminal {
return &openAIGPTVirtualTerminal{
Histories: history,
openAIKey: openAIKey,
client: resty.New(),
}
}

func buildPrompt(histories []History, command string) string {
var sb strings.Builder

Expand All @@ -79,7 +81,7 @@ func buildPrompt(histories []History, command string) string {
return sb.String()
}

func (openAIGPTVirtualTerminal *OpenAIGPTVirtualTerminal) GetCompletions(command string) (string, error) {
func (openAIGPTVirtualTerminal *openAIGPTVirtualTerminal) GetCompletions(command string) (string, error) {
requestJson, err := json.Marshal(gptRequest{
Model: "text-davinci-003",
Prompt: buildPrompt(openAIGPTVirtualTerminal.Histories, command),
Expand All @@ -94,14 +96,14 @@ func (openAIGPTVirtualTerminal *OpenAIGPTVirtualTerminal) GetCompletions(command
return "", err
}

if openAIGPTVirtualTerminal.OpenAPIChatGPTSecretKey == "" {
return "", errors.New("OpenAPIChatGPTSecretKey is empty")
if openAIGPTVirtualTerminal.openAIKey == "" {
return "", errors.New("openAIKey is empty")
}

response, err := openAIGPTVirtualTerminal.client.R().
SetHeader("Content-Type", "application/json").
SetBody(requestJson).
SetAuthToken(openAIGPTVirtualTerminal.OpenAPIChatGPTSecretKey).
SetAuthToken(openAIGPTVirtualTerminal.openAIKey).
SetResult(&gptResponse{}).
Post(openAIGPTEndpoint)

Expand Down
44 changes: 39 additions & 5 deletions plugins/openai-gpt_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,15 @@ func TestBuildPromptWithHistory(t *testing.T) {
prompt)
}

func TestBuildGetCompletions(t *testing.T) {
func TestBuildGetCompletionsFailValidation(t *testing.T) {
openAIGPTVirtualTerminal := Init(make([]History, 0), "")

_, err := openAIGPTVirtualTerminal.GetCompletions("test")

assert.Equal(t, "openAIKey is empty", err.Error())
}

func TestBuildGetCompletionsWithResults(t *testing.T) {
client := resty.New()
httpmock.ActivateNonDefault(client.GetClient())
defer httpmock.DeactivateAndReset()
Expand All @@ -68,10 +76,8 @@ func TestBuildGetCompletions(t *testing.T) {
},
)

openAIGPTVirtualTerminal := OpenAIGPTVirtualTerminal{
OpenAPIChatGPTSecretKey: "sdjdnklfjndslkjanfk",
client: client,
}
openAIGPTVirtualTerminal := Init(make([]History, 0), "sdjdnklfjndslkjanfk")
openAIGPTVirtualTerminal.client = client

//When
str, err := openAIGPTVirtualTerminal.GetCompletions("ls")
Expand All @@ -80,3 +86,31 @@ func TestBuildGetCompletions(t *testing.T) {
assert.Nil(t, err)
assert.Equal(t, "prova.txt", str)
}

func TestBuildGetCompletionsWithoutResults(t *testing.T) {
client := resty.New()
httpmock.ActivateNonDefault(client.GetClient())
defer httpmock.DeactivateAndReset()

// Given
httpmock.RegisterResponder("POST", openAIGPTEndpoint,
func(req *http.Request) (*http.Response, error) {
resp, err := httpmock.NewJsonResponse(200, &gptResponse{
Choices: []Choice{},
})
if err != nil {
return httpmock.NewStringResponse(500, ""), nil
}
return resp, nil
},
)

openAIGPTVirtualTerminal := Init(make([]History, 0), "sdjdnklfjndslkjanfk")
openAIGPTVirtualTerminal.client = client

//When
_, err := openAIGPTVirtualTerminal.GetCompletions("ls")

//Then
assert.Equal(t, "no choices", err.Error())
}
2 changes: 1 addition & 1 deletion protocols/protocol_manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ type ProtocolManager struct {

func InitProtocolManager(tracerStrategy tracer.Strategy, strategy ServiceStrategy) *ProtocolManager {
return &ProtocolManager{
tracer: tracer.Init(tracerStrategy),
tracer: tracer.GetInstance(tracerStrategy),
strategy: strategy,
}
}
Expand Down
5 changes: 2 additions & 3 deletions protocols/strategies/ssh.go
Original file line number Diff line number Diff line change
Expand Up @@ -63,8 +63,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
commandOutput := command.Handler

if command.Plugin == plugins.ChatGPTPluginName {
openAIGPTVirtualTerminal := plugins.OpenAIGPTVirtualTerminal{Histories: histories, OpenAPIChatGPTSecretKey: beelzebubServiceConfiguration.Plugin.OpenAPIChatGPTSecretKey}
openAIGPTVirtualTerminal.InjectDependency()
openAIGPTVirtualTerminal := plugins.Init(histories, beelzebubServiceConfiguration.Plugin.OpenAPIChatGPTSecretKey)

if commandOutput, err = openAIGPTVirtualTerminal.GetCompletions(commandInput); err != nil {
log.Errorf("Error GetCompletions: %s, %s", commandInput, err.Error())
Expand Down Expand Up @@ -125,7 +124,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
log.WithFields(log.Fields{
"port": beelzebubServiceConfiguration.Address,
"commands": len(beelzebubServiceConfiguration.Commands),
}).Infof("Init service %s", beelzebubServiceConfiguration.Protocol)
}).Infof("GetInstance service %s", beelzebubServiceConfiguration.Protocol)
return nil
}

Expand Down

0 comments on commit 07ffdd8

Please sign in to comment.