Repository: clouditor/clouditor
Pull request #700 (2022-05-22 20:08:44)

Using `api.StreamsOf` in discovery service
This PR removes the custom streams implementation from the discovery service and migrates it to `api.StreamsOf`, similar to the other services.

codecov[bot]: # [Codecov](https://codecov.io/gh/clouditor/clouditor/pull/700?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=clouditor) Report
> Merging [#700](https://codecov.io/gh/clouditor/clouditor/pull/700?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=clouditor) (34fb062) into [main](https://codecov.io/gh/clouditor/clouditor/commit/275f86ebaf140b1fffa15a1d8a8153ce16683b27?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=clouditor) (275f86e) will **decrease** coverage by `0.08%`.
> The diff coverage is `96.42%`.

```diff
@@            Coverage Diff             @@
##             main     #700      +/-   ##
==========================================
- Coverage   81.48%   81.40%   -0.09%     
==========================================
  Files          65       65              
  Lines        6072     6065       -7     
==========================================
- Hits         4948     4937      -11     
- Misses        865      871       +6     
+ Partials      259      257       -2     
```

| Flag | Coverage Δ | |
|---|---|---|
| unittests | `81.40% <96.42%> (-0.09%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=clouditor#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/clouditor/clouditor/pull/700?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=clouditor) | Coverage Δ | |
|---|---|---|
| [service/discovery/discovery.go](https://codecov.io/gh/clouditor/clouditor/pull/700/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=clouditor#diff-c2VydmljZS9kaXNjb3ZlcnkvZGlzY292ZXJ5Lmdv) | `92.75% <96.42%> (+2.47%)` | :arrow_up: |
| [cli/commands/service/discovery/discovery.go](https://codecov.io/gh/clouditor/clouditor/pull/700/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=clouditor#diff-Y2xpL2NvbW1hbmRzL3NlcnZpY2UvZGlzY292ZXJ5L2Rpc2NvdmVyeS5nbw==) | `64.70% <0.00%> (-19.12%)` | :arrow_down: |
| [service/evidence/evidence\_store.go](https://codecov.io/gh/clouditor/clouditor/pull/700/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=clouditor#diff-c2VydmljZS9ldmlkZW5jZS9ldmlkZW5jZV9zdG9yZS5nbw==) | `93.47% <0.00%> (+0.14%)` | :arrow_up: |
| [api/streams.go](https://codecov.io/gh/clouditor/clouditor/pull/700/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=clouditor#diff-YXBpL3N0cmVhbXMuZ28=) | `95.57% <0.00%> (+2.65%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/clouditor/clouditor/pull/700?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=clouditor).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=clouditor)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/clouditor/clouditor/pull/700?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=clouditor). Last update [275f86e...34fb062](https://codecov.io/gh/clouditor/clouditor/pull/700?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=clouditor). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=clouditor).


Patch:
diff --git a/.github/labeler.yml b/.github/labeler.yml
index 280b41c4..f4abad35 100644
--- a/.github/labeler.yml
+++ b/.github/labeler.yml
@@ -13,6 +13,9 @@
 "api":
   - any: ["api/**/*", "!api/evaluation/*"]
 
+"service/discovery":
+  - service/discovery/*
+
 "service/assessment":
   - service/assessment/*
 
@@ -20,7 +23,13 @@
   - service/orchestrator/*
 
 "service":
-  - any: ["service/**/*", "!service/assessment/*", "!service/orchestrator/*"]
+  - any:
+      [
+        "service/**/*",
+        "!service/assessment/*",
+        "!service/orchestrator/*",
+        "!service/discovery/*",
+      ]
 
 "internal":
   - internal/**/*
diff --git a/service/discovery/discovery.go b/service/discovery/discovery.go
index 62de972f..5e4b4528 100644
--- a/service/discovery/discovery.go
+++ b/service/discovery/discovery.go
@@ -29,7 +29,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"io"
 	"sort"
 	"time"
 
@@ -64,6 +63,11 @@ const (
 
 var log *logrus.Entry
 
+type grpcTarget struct {
+	target string
+	opts   []grpc.DialOption
+}
+
 // Service is an implementation of the Clouditor Discovery service.
 // It should not be used directly, but rather the NewService constructor
 // should be used.
@@ -72,8 +76,8 @@ type Service struct {
 
 	configurations map[discovery.Discoverer]*Configuration
 
-	assessmentStream  assessment.Assessment_AssessEvidencesClient
-	assessmentAddress string
+	assessmentStreams *api.StreamsOf[assessment.Assessment_AssessEvidencesClient, *assessment.AssessEvidenceRequest]
+	assessmentAddress grpcTarget
 
 	resources map[string]voc.IsCloudResource
 	scheduler *gocron.Scheduler
@@ -100,9 +104,12 @@ const (
 type ServiceOption func(*Service)
 
 // WithAssessmentAddress is an option to configure the assessment service gRPC address.
-func WithAssessmentAddress(address string) ServiceOption {
+func WithAssessmentAddress(address string, opts ...grpc.DialOption) ServiceOption {
 	return func(s *Service) {
-		s.assessmentAddress = address
+		s.assessmentAddress = grpcTarget{
+			target: address,
+			opts:   opts,
+		}
 	}
 }
 
@@ -127,10 +134,13 @@ func WithProviders(providersList []string) ServiceOption {
 
 func NewService(opts ...ServiceOption) *Service {
 	s := &Service{
-		assessmentAddress: DefaultAssessmentAddress,
-		resources:         make(map[string]voc.IsCloudResource),
-		scheduler:         gocron.NewScheduler(time.UTC),
-		configurations:    make(map[discovery.Discoverer]*Configuration),
+		assessmentStreams: api.NewStreamsOf(api.WithLogger[assessment.Assessment_AssessEvidencesClient, *assessment.AssessEvidenceRequest](log)),
+		assessmentAddress: grpcTarget{
+			target: DefaultAssessmentAddress,
+		},
+		resources:      make(map[string]voc.IsCloudResource),
+		scheduler:      gocron.NewScheduler(time.UTC),
+		configurations: make(map[discovery.Discoverer]*Configuration),
 	}
 
 	// Apply any options
@@ -153,46 +163,29 @@ func (s *Service) Authorizer() api.Authorizer {
 
 // initAssessmentStream initializes the stream that is used to send evidences to the assessment service.
 // If configured, it uses the Authorizer of the discovery service to authenticate requests to the assessment.
-func (s *Service) initAssessmentStream(additionalOpts ...grpc.DialOption) error {
-	log.Infof("Trying to establish a connection to assessment service @ %v", s.assessmentAddress)
+func (s *Service) initAssessmentStream(target string, additionalOpts ...grpc.DialOption) (stream assessment.Assessment_AssessEvidencesClient, err error) {
+	log.Infof("Trying to establish a connection to assessment service @ %v", target)
 
 	// Establish connection to assessment gRPC service
-	conn, err := grpc.Dial(s.assessmentAddress,
-		api.DefaultGrpcDialOptions(s.assessmentAddress, s, additionalOpts...)...,
+	conn, err := grpc.Dial(target,
+		api.DefaultGrpcDialOptions(target, s, additionalOpts...)...,
 	)
 	if err != nil {
-		return fmt.Errorf("could not connect to assessment service: %w", err)
+		return nil, fmt.Errorf("could not connect to assessment service: %w", err)
 	}
 
 	client := assessment.NewAssessmentClient(conn)
 
 	// Set up the stream and store it in our service struct, so we can access it later to actually
 	// send the evidence data
-	s.assessmentStream, err = client.AssessEvidences(context.Background())
+	stream, err = client.AssessEvidences(context.Background())
 	if err != nil {
-		return fmt.Errorf("could not set up stream for assessing evidences: %w", err)
+		return nil, fmt.Errorf("could not set up stream for assessing evidences: %w", err)
 	}
 
-	// Receive responses from Assessment
-	// Currently we do not process the responses
-	go func() {
-		for {
-			_, err := s.assessmentStream.Recv()
-			if errors.Is(err, io.EOF) {
-				log.Debugf("no more requests in assessment stream available")
-				break
-			}
-
-			if err != nil {
-				log.Errorf("error receiving response from assessment stream: %+v", err)
-				break
-			}
-		}
-	}()
-
 	log.Infof("Connected to Assessment")
 
-	return nil
+	return
 }
 
 // Start starts discovery
@@ -202,15 +195,6 @@ func (s *Service) Start(_ context.Context, _ *discovery.StartDiscoveryRequest) (
 	log.Infof("Starting discovery...")
 	s.scheduler.TagsUnique()
 
-	// Establish connection to assessment component
-	if s.assessmentStream == nil {
-		err = s.initAssessmentStream()
-		if err != nil {
-			log.Errorf("Could not initialize stream to Assessment: %v", err)
-			return nil, status.Error(codes.Internal, "could not initialize stream to Assessment")
-		}
-	}
-
 	var discoverer []discovery.Discoverer
 
 	// Configure discoverers for given providers
@@ -279,7 +263,7 @@ func (s *Service) Shutdown() {
 	s.scheduler.Stop()
 }
 
-func (s *Service) StartDiscovery(discoverer discovery.Discoverer) {
+func (svc *Service) StartDiscovery(discoverer discovery.Discoverer) {
 	var (
 		err  error
 		list []voc.IsCloudResource
@@ -293,7 +277,7 @@ func (s *Service) StartDiscovery(discoverer discovery.Discoverer) {
 	}
 
 	for _, resource := range list {
-		s.resources[string(resource.GetID())] = resource
+		svc.resources[string(resource.GetID())] = resource
 
 		var (
 			v *structpb.Value
@@ -313,14 +297,15 @@ func (s *Service) StartDiscovery(discoverer discovery.Discoverer) {
 			Resource:  v,
 		}
 
-		if s.assessmentStream == nil {
-			log.Warnf("Evidence stream to Assessment component not available")
+		// Get Evidence Store stream
+		channel, err := svc.assessmentStreams.GetStream(svc.assessmentAddress.target, "Assessment", svc.initAssessmentStream, svc.assessmentAddress.opts...)
+		if err != nil {
+			err = fmt.Errorf("could not get stream to assessment service (%s): %w", svc.assessmentAddress.target, err)
+			log.Error(err)
 			continue
 		}
 
-		if err = s.assessmentStream.Send(&assessment.AssessEvidenceRequest{Evidence: e}); err != nil {
-			log.WithError(handleError(err, "Assessment"))
-		}
+		channel.Send(&assessment.AssessEvidenceRequest{Evidence: e})
 	}
 }
 


Test patch:
diff --git a/cli/commands/service/discovery/discovery_test.go b/cli/commands/service/discovery/discovery_test.go
index 5396e686..630ad45f 100644
--- a/cli/commands/service/discovery/discovery_test.go
+++ b/cli/commands/service/discovery/discovery_test.go
@@ -65,21 +65,6 @@ func TestAddCommands(t *testing.T) {
 	t.Errorf("No query command was added")
 }
 
-// TestNewStartDiscoveryCommand currently tests only the error case
-// For more tests mock stream to assessment
-func TestNewStartDiscoveryCommand(t *testing.T) {
-	var err error
-
-	cmd := NewStartDiscoveryCommand()
-	err = cmd.RunE(nil, []string{})
-	assert.Error(t, err)
-
-	var response = &discovery.StartDiscoveryResponse{}
-	assert.NotNil(t, response)
-	assert.False(t, response.Successful)
-
-}
-
 func TestNewQueryDiscoveryCommand(t *testing.T) {
 	var err error
 	var b bytes.Buffer
diff --git a/service/discovery/discovery_test.go b/service/discovery/discovery_test.go
index b4532237..42e25453 100644
--- a/service/discovery/discovery_test.go
+++ b/service/discovery/discovery_test.go
@@ -32,20 +32,19 @@ import (
 	"io"
 	"os"
 	"reflect"
+	"sync"
 	"testing"
 	"time"
 
+	"clouditor.io/clouditor/api"
 	"clouditor.io/clouditor/api/assessment"
 	"clouditor.io/clouditor/api/discovery"
 	"clouditor.io/clouditor/api/evidence"
-	"clouditor.io/clouditor/internal/testutil"
 	"clouditor.io/clouditor/internal/testutil/clitest"
 	"clouditor.io/clouditor/service/discovery/azure"
 	"clouditor.io/clouditor/voc"
 
-	oauth2 "github.com/oxisto/oauth2go"
 	"github.com/stretchr/testify/assert"
-	"golang.org/x/oauth2/clientcredentials"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
@@ -81,7 +80,8 @@ func TestNewService(t *testing.T) {
 				},
 			},
 			want: &Service{
-				assessmentAddress: "localhost:9091",
+				assessmentStreams: api.NewStreamsOf(api.WithLogger[assessment.Assessment_AssessEvidencesClient, *assessment.AssessEvidenceRequest](log)),
+				assessmentAddress: grpcTarget{target: "localhost:9091"},
 				resources:         make(map[string]voc.IsCloudResource),
 				configurations:    make(map[discovery.Discoverer]*Configuration),
 			},
@@ -104,12 +104,8 @@ func TestNewService(t *testing.T) {
 }
 
 func TestStartDiscovery(t *testing.T) {
-	discoveryService := NewService()
-
 	type fields struct {
-		assessmentStream    assessment.Assessment_AssessEvidencesClient
-		evidenceStoreStream evidence.EvidenceStore_StoreEvidencesClient
-		discoverer          discovery.Discoverer
+		discoverer discovery.Discoverer
 	}
 
 	tests := []struct {
@@ -120,65 +116,65 @@ func TestStartDiscovery(t *testing.T) {
 		{
 			name: "Err in discoverer",
 			fields: fields{
-				discoverer: mockDiscoverer{testCase: 0}},
+				discoverer: mockDiscoverer{testCase: 0},
+			},
 		},
 		{
 			name: "Err in marshaling the resource containing circular dependencies",
 			fields: fields{
-				discoverer: mockDiscoverer{testCase: 1}},
-		},
-		{
-			name: "No err in discoverer but no evidence stream to assessment",
-			fields: fields{
-				discoverer: mockDiscoverer{testCase: 2},
+				discoverer: mockDiscoverer{testCase: 1},
 			},
 		},
-		{
-			name: "No err in discoverer but no evidence stream to assessment available",
-			fields: fields{
-				assessmentStream: &mockAssessmentStream{},
-				discoverer:       mockDiscoverer{testCase: 2}},
-		},
-		{
-			name: "No err in discoverer but streaming to assessment fails",
-			fields: fields{
-				assessmentStream:    &mockAssessmentStream{},
-				evidenceStoreStream: mockEvidenceStoreStream{},
-				discoverer:          mockDiscoverer{testCase: 2}},
-		},
 		{
 			name: "No err",
 			fields: fields{
-				assessmentStream:    &mockAssessmentStream{connectionEstablished: true},
-				evidenceStoreStream: mockEvidenceStoreStream{},
-				discoverer:          mockDiscoverer{testCase: 2}},
+				discoverer: mockDiscoverer{testCase: 2},
+			},
 			checkEvidence: true,
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			discoveryService.assessmentStream = tt.fields.assessmentStream
-			discoveryService.StartDiscovery(tt.fields.discoverer)
+			mockStream := &mockAssessmentStream{connectionEstablished: true, expected: 2}
+			mockStream.Prepare()
+
+			svc := NewService()
+			svc.assessmentStreams = api.NewStreamsOf[assessment.Assessment_AssessEvidencesClient, *assessment.AssessEvidenceRequest]()
+			_, _ = svc.assessmentStreams.GetStream("mock", "Assessment", func(target string, additionalOpts ...grpc.DialOption) (stream assessment.Assessment_AssessEvidencesClient, err error) {
+				return mockStream, nil
+			})
+			svc.assessmentAddress = grpcTarget{target: "mock"}
+			go svc.StartDiscovery(tt.fields.discoverer)
 
-			// APIs for assessment and evidence store both send the same evidence. Thus, testing one is enough.
 			if tt.checkEvidence {
-				e := discoveryService.assessmentStream.(*mockAssessmentStream).sentEvidence
+				mockStream.Wait()
+
+				want, _ := tt.fields.discoverer.List()
+
+				got := mockStream.sentEvidences
+				assert.Equal(t, len(want), len(got))
+
+				// Retrieve the last one
+				e := got[len(got)-1]
+
+				// Check, if evidence was sent
+				assert.NotNil(t, e)
 				// Check if UUID has been created
 				assert.NotEmpty(t, e.Id)
 				// Check if cloud resources / properties are there
 				assert.NotEmpty(t, e.Resource)
 				// Check if ID of mockDiscovery's resource is mapped to resource id of the evidence
-				list, _ := tt.fields.discoverer.List()
+
 				// Only the last element sent can be checked
-				assert.Equal(t, string(list[len(list)-1].GetID()), e.Resource.GetStructValue().AsMap()["id"].(string))
+				assert.Equal(t, string(want[len(want)-1].GetID()), e.Resource.GetStructValue().AsMap()["id"].(string))
 			}
 		})
 	}
 }
 
 func TestQuery(t *testing.T) {
-	s := NewService()
+	s := NewService(WithAssessmentAddress("bufnet", grpc.WithContextDialer(bufConnDialer)))
 	s.StartDiscovery(mockDiscoverer{testCase: 2})
 
 	type fields struct {
@@ -243,8 +239,7 @@ func TestStart(t *testing.T) {
 	}
 
 	type fields struct {
-		hasRPCConnection bool
-		envVariables     []envVariable
+		envVariables []envVariable
 	}
 
 	tests := []struct {
@@ -259,7 +254,6 @@ func TestStart(t *testing.T) {
 		{
 			name: "Azure authorizer from ENV",
 			fields: fields{
-				hasRPCConnection: true,
 				envVariables: []envVariable{
 					// We must set AZURE_AUTH_LOCATION to the Azure credentials test file and the set HOME to a
 					// wrong path so that the Azure authorizer passes and the K8S authorizer fails
@@ -288,7 +282,6 @@ func TestStart(t *testing.T) {
 		{
 			name: "Azure authorizer from file",
 			fields: fields{
-				hasRPCConnection: true,
 				envVariables: []envVariable{
 					// We must set AZURE_AUTH_LOCATION to the Azure credentials test file and the set HOME to a
 					// wrong path so that the Azure authorizer passes and the K8S authorizer fails
@@ -313,7 +306,6 @@ func TestStart(t *testing.T) {
 		{
 			name: "No Azure authorizer",
 			fields: fields{
-				hasRPCConnection: true,
 				// We must set env variables accordingly s.t. all authorizer will fail
 				envVariables: []envVariable{
 					// Let `authorizer from ENV` fail
@@ -379,7 +371,6 @@ func TestStart(t *testing.T) {
 		{
 			name: "No K8s authorizer",
 			fields: fields{
-				hasRPCConnection: true,
 				envVariables: []envVariable{
 					// We must set HOME to a wrong path so that the K8S authorizer fails
 					{
@@ -394,20 +385,9 @@ func TestStart(t *testing.T) {
 			wantErr:        true,
 			wantErrMessage: "could not authenticate to Kubernetes",
 		},
-		{
-			name: "No RPC connection",
-			fields: fields{
-				hasRPCConnection: false,
-			},
-			providers:      []string{"aws", "azure", "k8s"},
-			wantResp:       nil,
-			wantErr:        true,
-			wantErrMessage: "could not initialize stream to Assessment",
-		},
 		{
 			name: "Request with 2 providers",
 			fields: fields{
-				hasRPCConnection: true,
 				envVariables: []envVariable{
 					// We must set HOME to a wrong path so that the AWS and k8s authorizer fails in all systems, regardless if AWS and k8s paths are set or not
 					{
@@ -423,20 +403,16 @@ func TestStart(t *testing.T) {
 			wantErrMessage: "could not authenticate to",
 		},
 		{
-			name: "Empty request",
-			fields: fields{
-				hasRPCConnection: true,
-			},
+			name:           "Empty request",
+			fields:         fields{},
 			providers:      []string{},
 			wantResp:       &discovery.StartDiscoveryResponse{Successful: true},
 			wantErr:        false,
 			wantErrMessage: "",
 		},
 		{
-			name: "Request with wrong provider name",
-			fields: fields{
-				hasRPCConnection: true,
-			},
+			name:           "Request with wrong provider name",
+			fields:         fields{},
 			providers:      []string{"falseProvider"},
 			wantResp:       nil,
 			wantErr:        true,
@@ -447,9 +423,6 @@ func TestStart(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			s := NewService(WithProviders(tt.providers))
-			if tt.fields.hasRPCConnection {
-				assert.NoError(t, s.initAssessmentStream(grpc.WithContextDialer(bufConnDialer)))
-			}
 
 			for _, env := range tt.fields.envVariables {
 				if env.hasEnvVariable {
@@ -472,96 +445,6 @@ func TestStart(t *testing.T) {
 	}
 }
 
-func TestService_initAssessmentStream(t *testing.T) {
-	var (
-		authSrv *oauth2.AuthorizationServer
-		port    int
-		err     error
-	)
-
-	authSrv, port, err = testutil.StartAuthenticationServer()
-	defer func() {
-		err = authSrv.Close()
-		assert.NoError(t, err)
-	}()
-
-	assert.NoError(t, err)
-
-	type fields struct {
-		hasRPCConnection bool
-		clientID         string
-		clientSecret     string
-	}
-
-	tests := []struct {
-		name    string
-		fields  fields
-		wantErr assert.ErrorAssertionFunc
-	}{
-		{
-			name: "No RPC connection",
-			fields: fields{
-				hasRPCConnection: false,
-			},
-			wantErr: func(tt assert.TestingT, err error, i ...interface{}) bool {
-				// We are looking for a connection refused message
-				innerErr := errors.Unwrap(err)
-				s, _ := status.FromError(innerErr)
-
-				if s.Code() != codes.Unavailable {
-					tt.Errorf("Status should be codes.Unavailable: %v", s.Code())
-					return false
-				}
-
-				return true
-			},
-		},
-		{
-			name: "Authenticated RPC connection with valid client credentials",
-			fields: fields{
-				hasRPCConnection: true,
-				clientID:         testutil.TestAuthClientID,
-				clientSecret:     testutil.TestAuthClientSecret,
-			},
-			wantErr: nil,
-		},
-		{
-			name: "Authenticated RPC connection with invalid client credentials",
-			fields: fields{
-				hasRPCConnection: true,
-				clientID:         "not" + testutil.TestAuthClientID,
-				clientSecret:     testutil.TestAuthClientSecret,
-			},
-			wantErr: func(tt assert.TestingT, err error, i ...interface{}) bool {
-				s, _ := status.FromError(errors.Unwrap(err))
-				return assert.Equal(t, codes.Unauthenticated, s.Code())
-			},
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			s := NewService(WithOAuth2Authorizer(&clientcredentials.Config{
-				ClientID:     tt.fields.clientID,
-				ClientSecret: tt.fields.clientSecret,
-				TokenURL:     testutil.TokenURL(port),
-			}))
-
-			var opts []grpc.DialOption
-			if tt.fields.hasRPCConnection {
-				// Make this a valid RPC connection by connecting to our bufnet service
-				opts = []grpc.DialOption{grpc.WithContextDialer(bufConnDialer)}
-				s.assessmentAddress = "bufnet"
-			}
-
-			err := s.initAssessmentStream(opts...)
-			if tt.wantErr != nil {
-				tt.wantErr(t, err)
-			}
-		})
-	}
-}
-
 func TestShutdown(t *testing.T) {
 	service := NewService()
 	service.Shutdown()
@@ -674,10 +557,20 @@ func wrongFormattedResource() voc.IsCloudResource {
 // mockAssessmentStream implements Assessment_AssessEvidencesClient interface
 type mockAssessmentStream struct {
 	// We add sentEvidence field to test the evidence that would be sent over gRPC
-	sentEvidence *evidence.Evidence
+	sentEvidences []*evidence.Evidence
 	// We add connectionEstablished to differentiate between the case where evidences can be sent and not
 	connectionEstablished bool
 	counter               int
+	expected              int
+	wg                    sync.WaitGroup
+}
+
+func (m *mockAssessmentStream) Prepare() {
+	m.wg.Add(m.expected)
+}
+
+func (m *mockAssessmentStream) Wait() {
+	m.wg.Wait()
 }
 
 func (m *mockAssessmentStream) Recv() (*assessment.AssessEvidenceResponse, error) {
@@ -698,13 +591,7 @@ func (m *mockAssessmentStream) Recv() (*assessment.AssessEvidenceResponse, error
 }
 
 func (m *mockAssessmentStream) Send(req *assessment.AssessEvidenceRequest) (err error) {
-	e := req.Evidence
-	if m.connectionEstablished {
-		m.sentEvidence = e
-	} else {
-		err = fmt.Errorf("mock send error")
-	}
-	return
+	return m.SendMsg(req)
 }
 
 func (*mockAssessmentStream) CloseAndRecv() (*emptypb.Empty, error) {
@@ -727,51 +614,20 @@ func (*mockAssessmentStream) Context() context.Context {
 	return nil
 }
 
-func (*mockAssessmentStream) SendMsg(_ interface{}) error {
-	return nil
-}
-
-func (*mockAssessmentStream) RecvMsg(_ interface{}) error {
-	return nil
-}
-
-// mockEvidenceStoreStream implements EvidenceStore_StoreEvidencesClient interface
-type mockEvidenceStoreStream struct {
-}
-
-func (mockEvidenceStoreStream) Recv() (*evidence.StoreEvidenceResponse, error) {
-	return nil, nil
-}
-
-func (mockEvidenceStoreStream) Send(_ *evidence.StoreEvidenceRequest) error {
-	return fmt.Errorf("mock send error")
-}
-
-func (mockEvidenceStoreStream) CloseAndRecv() (*emptypb.Empty, error) {
-	return nil, nil
-}
+func (m *mockAssessmentStream) SendMsg(req interface{}) (err error) {
+	m.wg.Done()
 
-func (mockEvidenceStoreStream) Header() (metadata.MD, error) {
-	return nil, nil
-}
-
-func (mockEvidenceStoreStream) Trailer() metadata.MD {
-	return nil
-}
-
-func (mockEvidenceStoreStream) CloseSend() error {
-	return nil
-}
-
-func (mockEvidenceStoreStream) Context() context.Context {
-	return nil
-}
+	e := req.(*assessment.AssessEvidenceRequest).Evidence
+	if m.connectionEstablished {
+		m.sentEvidences = append(m.sentEvidences, e)
+	} else {
+		err = fmt.Errorf("mock send error")
+	}
 
-func (mockEvidenceStoreStream) SendMsg(_ interface{}) error {
-	return nil
+	return
 }
 
-func (mockEvidenceStoreStream) RecvMsg(_ interface{}) error {
+func (*mockAssessmentStream) RecvMsg(_ interface{}) error {
 	return nil
 }