Skip to content

Commit

Permalink
Fixed unit test
Browse files Browse the repository at this point in the history
  • Loading branch information
@SchawnnDev
SchawnnDev committed May 28, 2024
1 parent db41fb1 commit a45e105
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 35 deletions.
31 changes: 22 additions & 9 deletions internals/export/wrapper.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,16 +85,25 @@ type Wrapper struct {
}

// NewWrapperItem creates a new export wrapper item
func NewWrapperItem(facts []engine.Fact, title string, params CSVParameters, user users.User, factParameters map[string]string) *WrapperItem {
//
// addHashPrefix adds a hash as prefix to resulting files, to avoid duplicates
func NewWrapperItem(facts []engine.Fact, title string, params CSVParameters, user users.User, factParameters map[string]string, addHashPrefix bool) *WrapperItem {
var factIDs []int64
for _, fact := range facts {
factIDs = append(factIDs, fact.ID)
}

// file extension should be gz
// add random string to avoid multiple files with same name
fileName := security.RandStringWithCharset(5, randCharSet) + "_" +
strings.ReplaceAll(title, " ", "_") + ".csv.gz"
var fileName string
if addHashPrefix {
// file extension should be gz
// add random string to avoid multiple files with same name
fileName = security.RandStringWithCharset(5, randCharSet) + "_" +
strings.ReplaceAll(title, " ", "_") + ".csv.gz"
} else if !strings.HasSuffix(title, ".csv.gz") {
fileName = strings.ReplaceAll(title, " ", "_") + ".csv.gz"
} else {
fileName = strings.ReplaceAll(title, " ", "_")
}

return &WrapperItem{
Users: append([]string{}, user.Login),
Expand Down Expand Up @@ -188,7 +197,9 @@ func factsEquals(a, b []engine.Fact) bool {
}

// AddToQueue Adds a new export to the export worker queue
func (ew *Wrapper) AddToQueue(facts []engine.Fact, title string, params CSVParameters, user users.User, factParameters map[string]string) (*WrapperItem, int) {
//
// addHashPrefix adds a hash as prefix to resulting files, to avoid duplicates
func (ew *Wrapper) AddToQueue(facts []engine.Fact, title string, params CSVParameters, user users.User, factParameters map[string]string, addHashPrefix bool) (*WrapperItem, int) {
ew.queueMutex.Lock()
defer ew.queueMutex.Unlock()

Expand All @@ -212,13 +223,15 @@ func (ew *Wrapper) AddToQueue(facts []engine.Fact, title string, params CSVParam
return nil, CodeQueueFull
}

item := NewWrapperItem(facts, title, params, user, factParameters)
item := NewWrapperItem(facts, title, params, user, factParameters, addHashPrefix)
ew.queue = append(ew.queue, item)
return item, CodeAdded
}

// AddToQueueCustom Adds a new export with a custom elastic connection and a custom search request
func (ew *Wrapper) AddToQueueCustom(elasticName string, request *search.Request, indices, title string, params CSVParameters, user users.User) (*WrapperItem, int) {
//
// addHashPrefix adds a hash as prefix to resulting files, to avoid duplicates
func (ew *Wrapper) AddToQueueCustom(elasticName string, request *search.Request, indices, title string, params CSVParameters, user users.User, addHashPrefix bool) (*WrapperItem, int) {
ew.queueMutex.Lock()
defer ew.queueMutex.Unlock()

Expand All @@ -242,7 +255,7 @@ func (ew *Wrapper) AddToQueueCustom(elasticName string, request *search.Request,
return nil, CodeQueueFull
}

item := NewWrapperItem([]engine.Fact{}, title, params, user, map[string]string{})
item := NewWrapperItem([]engine.Fact{}, title, params, user, map[string]string{}, addHashPrefix)
item.Custom = true
item.ElasticName = elasticName
item.SearchRequest = request
Expand Down
51 changes: 25 additions & 26 deletions internals/export/wrapper_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ func TestFactsEquals(t *testing.T) {
}

func TestNewWrapperItem(t *testing.T) {
item := NewWrapperItem([]engine.Fact{{ID: 1}}, "test", CSVParameters{}, users.User{Login: "test"}, make(map[string]string))
item := NewWrapperItem([]engine.Fact{{ID: 1}}, "test", CSVParameters{}, users.User{Login: "test"}, make(map[string]string), false)
expression.AssertNotEqual(t, item.Id, "")
expression.AssertEqual(t, factsEquals(item.Facts, []engine.Fact{{ID: 1}}), true)
expression.AssertEqual(t, item.Params.Equals(CSVParameters{}), true)
Expand All @@ -45,7 +45,7 @@ func TestNewWrapperItem(t *testing.T) {
}

func TestWrapperItem_ContainsFact(t *testing.T) {
item := NewWrapperItem([]engine.Fact{{ID: 1}, {ID: 22}, {ID: 33}}, "test.txt", CSVParameters{}, users.User{Login: "test"}, make(map[string]string))
item := NewWrapperItem([]engine.Fact{{ID: 1}, {ID: 22}, {ID: 33}}, "test.txt", CSVParameters{}, users.User{Login: "test"}, make(map[string]string), false)
expression.AssertEqual(t, item.ContainsFact(1), true)
expression.AssertEqual(t, item.ContainsFact(22), true)
expression.AssertEqual(t, item.ContainsFact(3), false)
Expand All @@ -70,13 +70,13 @@ func TestAddToQueue(t *testing.T) {
user1 := users.User{Login: "bla"}
user2 := users.User{Login: "blabla"}
csvParams := CSVParameters{}
_, result := wrapper.AddToQueue([]engine.Fact{{ID: 1}}, "test.txt", csvParams, user1, make(map[string]string))
_, result := wrapper.AddToQueue([]engine.Fact{{ID: 1}}, "test.txt", csvParams, user1, make(map[string]string), false)
expression.AssertEqual(t, result, CodeAdded, "AddToQueue should return CodeAdded")
_, result = wrapper.AddToQueue([]engine.Fact{{ID: 1}}, "test.txt", csvParams, user1, make(map[string]string))
_, result = wrapper.AddToQueue([]engine.Fact{{ID: 1}}, "test.txt", csvParams, user1, make(map[string]string), false)
expression.AssertEqual(t, result, CodeUserExists, "AddToQueue should return CodeUserExists")
_, result = wrapper.AddToQueue([]engine.Fact{{ID: 1}}, "test.txt", csvParams, user2, make(map[string]string))
_, result = wrapper.AddToQueue([]engine.Fact{{ID: 1}}, "test.txt", csvParams, user2, make(map[string]string), false)
expression.AssertEqual(t, result, CodeUserAdded, "AddToQueue should return CodeUserAdded")
_, result = wrapper.AddToQueue([]engine.Fact{{ID: 2}}, "test.txt", csvParams, user2, make(map[string]string))
_, result = wrapper.AddToQueue([]engine.Fact{{ID: 2}}, "test.txt", csvParams, user2, make(map[string]string), false)
expression.AssertEqual(t, result, CodeQueueFull, "AddToQueue should return CodeQueueFull")
}

Expand All @@ -90,8 +90,7 @@ func TestStartDispatcher(t *testing.T) {
}
defer os.RemoveAll(dname)

// create a file that is 2 days old
file, err := os.CreateTemp(dname, "exportdispatcher")
file, err := os.Create(filepath.Join(dname, "exportdispatcher.csv.gz"))
if err != nil {
t.Error(err)
t.FailNow()
Expand Down Expand Up @@ -119,7 +118,7 @@ func TestStartDispatcher(t *testing.T) {

// add a task to the queue and check if the task was added to queue
user := users.User{Login: "test"}
_, result := wrapper.AddToQueue([]engine.Fact{{ID: 1}}, fileName, CSVParameters{}, user, make(map[string]string))
_, result := wrapper.AddToQueue([]engine.Fact{{ID: 1, Intent: &engine.IntentFragment{}}}, fileName, CSVParameters{}, user, make(map[string]string), false)
expression.AssertEqual(t, result, CodeAdded, "AddToQueue should return CodeAdded")
wrapper.queueMutex.Lock()
expression.AssertEqual(t, len(wrapper.queue), 1)
Expand Down Expand Up @@ -225,16 +224,16 @@ func TestWrapper_GetUserExports(t *testing.T) {
wrapper := NewWrapper("/tmp", 1, 1, 2)
user1 := users.User{Login: "bla"}
user2 := users.User{Login: "blabla"}
item1 := NewWrapperItem([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, user1, make(map[string]string))
item2 := NewWrapperItem([]engine.Fact{{ID: 2}}, "test.txt", CSVParameters{}, user1, make(map[string]string))
item3 := NewWrapperItem([]engine.Fact{{ID: 3}}, "test.txt", CSVParameters{}, user1, make(map[string]string))
item4 := NewWrapperItem([]engine.Fact{{ID: 4}}, "test.txt", CSVParameters{}, user2, make(map[string]string))
item1 := NewWrapperItem([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, user1, make(map[string]string), false)
item2 := NewWrapperItem([]engine.Fact{{ID: 2}}, "test.txt", CSVParameters{}, user1, make(map[string]string), false)
item3 := NewWrapperItem([]engine.Fact{{ID: 3}}, "test.txt", CSVParameters{}, user1, make(map[string]string), false)
item4 := NewWrapperItem([]engine.Fact{{ID: 4}}, "test.txt", CSVParameters{}, user2, make(map[string]string), false)
wrapper.archive.Store(item1.Id, *item1)
wrapper.archive.Store(item2.Id, *item2)
wrapper.archive.Store(item3.Id, *item3)
wrapper.archive.Store(item4.Id, *item4)
wrapper.AddToQueue([]engine.Fact{{ID: 5}}, "test.txt", CSVParameters{}, user1, make(map[string]string))
wrapper.AddToQueue([]engine.Fact{{ID: 6}}, "test.txt", CSVParameters{}, user2, make(map[string]string))
wrapper.AddToQueue([]engine.Fact{{ID: 5}}, "test.txt", CSVParameters{}, user1, make(map[string]string), false)
wrapper.AddToQueue([]engine.Fact{{ID: 6}}, "test.txt", CSVParameters{}, user2, make(map[string]string), false)
exports := wrapper.GetUserExports(user1)
expression.AssertEqual(t, len(exports), 4)
exports = wrapper.GetUserExports(user2)
Expand All @@ -246,8 +245,8 @@ func TestWrapper_DequeueWrapperItem(t *testing.T) {
i, ok := wrapper.dequeueWrapperItem(&WrapperItem{})
expression.AssertEqual(t, ok, false)
expression.AssertEqual(t, i, 0)
wrapper.AddToQueue([]engine.Fact{{ID: 5}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string))
wrapper.AddToQueue([]engine.Fact{{ID: 6}}, "test.txt", CSVParameters{}, users.User{Login: "blabla"}, make(map[string]string))
wrapper.AddToQueue([]engine.Fact{{ID: 5}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string), false)
wrapper.AddToQueue([]engine.Fact{{ID: 6}}, "test.txt", CSVParameters{}, users.User{Login: "blabla"}, make(map[string]string), false)

expression.AssertEqual(t, len(wrapper.queue), 2)
item1 := wrapper.queue[0]
Expand Down Expand Up @@ -299,7 +298,7 @@ func TestWrapper_dispatchExportQueue(t *testing.T) {
expression.AssertEqual(t, worker.IsAvailable(), true, "worker should still be available, because no items in queue")

// we add an item to the queue
wrapper.AddToQueue([]engine.Fact{{ID: 1}}, fileName, CSVParameters{}, users.User{Login: "test"}, make(map[string]string))
wrapper.AddToQueue([]engine.Fact{{ID: 1}}, fileName, CSVParameters{}, users.User{Login: "test"}, make(map[string]string), false)

// we test if dispatchExportQueue will not dispatch the item, no worker available
worker.SwapAvailable(false)
Expand Down Expand Up @@ -330,7 +329,7 @@ func TestWrapper_dispatchExportQueue(t *testing.T) {

func TestWrapper_FindArchive(t *testing.T) {
wrapper := NewWrapper("/tmp", 1, 1, 2)
item := NewWrapperItem([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string))
item := NewWrapperItem([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string), false)
wrapper.archive.Store(item.Id, *item)

// testing with non-existing item in archive
Expand All @@ -347,14 +346,14 @@ func TestWrapper_FindArchive(t *testing.T) {
}

func TestWrapper_ContainsUser(t *testing.T) {
item := NewWrapperItem([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string))
item := NewWrapperItem([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string), false)
expression.AssertEqual(t, item.ContainsUser(users.User{Login: "bla"}), true)
expression.AssertEqual(t, item.ContainsUser(users.User{Login: "blabla"}), false)
}

func TestWrapper_DeleteExport(t *testing.T) {
wrapper := NewWrapper("/tmp", 1, 1, 2)
item := NewWrapperItem([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string))
item := NewWrapperItem([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string), false)

// test archive
wrapper.archive.Store(item.Id, *item)
Expand All @@ -371,7 +370,7 @@ func TestWrapper_DeleteExport(t *testing.T) {
item.Users = []string{"bla"}

// test queue
queueItem, code := wrapper.AddToQueue([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string))
queueItem, code := wrapper.AddToQueue([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string), false)
expression.AssertEqual(t, code, CodeAdded, "item should have been added to queue")
wrapper.queueMutex.Lock()
expression.AssertEqual(t, len(wrapper.queue), 1, "item should be in queue")
Expand All @@ -382,9 +381,9 @@ func TestWrapper_DeleteExport(t *testing.T) {
wrapper.queueMutex.Unlock()

// test queue multi-user
queueItem, code = wrapper.AddToQueue([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string))
queueItem, code = wrapper.AddToQueue([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string), false)
expression.AssertEqual(t, code, CodeAdded, "item should have been added to queue")
_, code = wrapper.AddToQueue([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "blabla"}, make(map[string]string))
_, code = wrapper.AddToQueue([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "blabla"}, make(map[string]string), false)
expression.AssertEqual(t, code, CodeUserAdded, "user should have been added to existing item in queue")
wrapper.queueMutex.Lock()
expression.AssertEqual(t, len(wrapper.queue), 1, "item should be in queue")
Expand Down Expand Up @@ -420,7 +419,7 @@ func TestWrapper_DeleteExport(t *testing.T) {

func TestWrapper_GetUserExport(t *testing.T) {
wrapper := NewWrapper("/tmp", 1, 1, 2)
item := NewWrapperItem([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string))
item := NewWrapperItem([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string), false)

// test item in archive
wrapper.archive.Store(item.Id, *item)
Expand All @@ -432,7 +431,7 @@ func TestWrapper_GetUserExport(t *testing.T) {
wrapper.archive.Delete(item.Id)

// test item in queue queue
queueItem, code := wrapper.AddToQueue([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string))
queueItem, code := wrapper.AddToQueue([]engine.Fact{{ID: 1}}, "test.txt", CSVParameters{}, users.User{Login: "bla"}, make(map[string]string), false)
expression.AssertEqual(t, code, CodeAdded, "item should have been added to queue")
export, ok = wrapper.GetUserExport(queueItem.Id, users.User{Login: "bla"})
expression.AssertEqual(t, ok, true)
Expand Down

0 comments on commit a45e105

Please sign in to comment.