package scanner
import (
"sync"
"testing"
"time"
"github.com/safedep/dry/adapters"
"github.com/stretchr/testify/assert"
"google.golang.org/grpc"
)
func TestMalysisMalwareEnricher(t *testing.T) {
t.Run("NewMalysisMalwareEnricher", func(t *testing.T) {
t.Run("should reject nil grpc client connection", func(t *testing.T) {
_, err := NewMalysisMalwareEnricher(nil, nil, DefaultMalysisMalwareEnricherConfig())
assert.Error(t, err)
assert.ErrorContains(t, err, "grpc client connection is required")
})
t.Run("should wait till max timeout", func(t *testing.T) {
enricher, err := NewMalysisMalwareEnricher(&grpc.ClientConn{}, &adapters.GithubClient{}, MalysisMalwareEnricherConfig{
Timeout: 1,
QueryWorkerCount: 0,
})
if err != nil {
t.Fatalf("failed to create enricher: %v", err)
}
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
_ = enricher.Wait()
}()
closer := make(chan bool)
go func() {
wg.Wait()
close(closer)
}()
timer := time.NewTimer(2 * time.Second)
select {
case <-timer.C:
assert.Fail(t, "timeout occurred instead of completion")
case <-closer:
timer.Stop()
}
})
})
}
func TestMalysisMalwareEnricherQueryWorker(t *testing.T) {
t.Run("MalysisMalwareEnricher/QueryWorker", func(t *testing.T) {
t.Run("should return error for requests with future retry", func(t *testing.T) {
enricher, err := NewMalysisMalwareEnricher(&grpc.ClientConn{}, &adapters.GithubClient{}, MalysisMalwareEnricherConfig{
Timeout: 5 * time.Second,
QueryWorkerCount: 1,
MaxQueryRetries: 3,
})
assert.NoError(t, err)
assert.NotNil(t, enricher)
assert.NotNil(t, enricher.queryChannel)
assert.NotNil(t, enricher.resultsChannel)
enricher.queryChannel <- &analysisQueryRequest{
analysisId: "test",
retryCount: 1,
nextRetryAt: time.Now().Add(5 * time.Second),
}
enricher.wg.Add(1)
waiterCh := make(chan bool)
go func() {
err = enricher.Wait()
assert.NoError(t, err, "error occurred while waiting for enricher")
close(waiterCh)
}()
select {
case result := <-enricher.resultsChannel:
assert.Error(t, result.err)
assert.ErrorContains(t, result.err, errMalysisPollRetryInFuture.Error())
case <-waiterCh:
assert.Fail(t, "timeout occurred instead of completion")
}
})
})
}
