github.com/google/fleetspeak@v0.1.15-0.20240426164851-4f31f62c1aea/fleetspeak/src/client/clitesting/mock_service_context.go (about)

     1  // Copyright 2017 Google Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     https://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package clitesting contains utilities useful for testing clients and client
    16  // components.
    17  package clitesting
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"time"
    23  
    24  	"github.com/google/fleetspeak/fleetspeak/src/client/service"
    25  	"github.com/google/fleetspeak/fleetspeak/src/client/stats"
    26  
    27  	fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
    28  )
    29  
    30  const (
    31  	// MockCommTimeout is how long we wait for the test to accept messages
    32  	// sent through MockServiceContext.
    33  	MockCommTimeout = 5 * time.Second
    34  )
    35  
    36  // MockServiceContext is a mock service.Context which allows a test to run a service
    37  // without a full Fleetspeak client.
    38  type MockServiceContext struct {
    39  	service.Context
    40  
    41  	OutChan        chan *fspb.Message
    42  	StatsCollector stats.Collector
    43  }
    44  
    45  // Send implements service.Context by copying the message to sc.OutChan.
    46  func (sc *MockServiceContext) Send(ctx context.Context, m service.AckMessage) error {
    47  	select {
    48  	case <-ctx.Done():
    49  		return ctx.Err()
    50  	case <-time.After(MockCommTimeout):
    51  		return fmt.Errorf("Timed out while waiting for the test to pick up output: %+v", m)
    52  	case sc.OutChan <- m.M:
    53  		if m.Ack != nil {
    54  			m.Ack()
    55  		}
    56  		return nil
    57  	}
    58  }
    59  
    60  // Stats implements service.Context by returning the assigned stats.Collector or a NoopCollector
    61  // if none was set.
    62  func (sc *MockServiceContext) Stats() stats.Collector {
    63  	if sc.StatsCollector == nil {
    64  		return stats.NoopCollector{}
    65  	}
    66  	return sc.StatsCollector
    67  }