github.com/choria-io/go-choria@v0.28.1-0.20240416190746-b3bf9c7d5a45/integration/agentharness/agent.go (about)

     1  // Copyright (c) 2022, R.I. Pienaar and the Choria Project contributors
     2  //
     3  // SPDX-License-Identifier: Apache-2.0
     4  
     5  // Package agentharness is a testing framework to mock any RPC agent based on it's DDL
     6  //
     7  // All actions declared in the DDL will be mocked using gomock, expectations can be
     8  // set using the Stub() method
     9  //
    10  // See the harness test integration suite for an example
    11  package agentharness
    12  
    13  import (
    14  	"context"
    15  	"fmt"
    16  
    17  	"github.com/choria-io/go-choria/inter"
    18  	"github.com/choria-io/go-choria/providers/agent/mcorpc"
    19  	addl "github.com/choria-io/go-choria/providers/agent/mcorpc/ddl/agent"
    20  	"github.com/golang/mock/gomock"
    21  	"github.com/onsi/ginkgo/v2"
    22  	"github.com/sirupsen/logrus"
    23  )
    24  
    25  // NewWithDDLBytes creates a new test harness based on a ddl contained in ddlBytes
    26  func NewWithDDLBytes(fw inter.Framework, ctl *gomock.Controller, name string, ddlBytes []byte) (*AgentHarness, error) {
    27  	ddl, err := addl.NewFromBytes(ddlBytes)
    28  	if err != nil {
    29  		return nil, err
    30  	}
    31  
    32  	return New(fw, ctl, name, ddl)
    33  }
    34  
    35  // NewWithDDLFile creates a new test harness based on a ddl contained on disk
    36  func NewWithDDLFile(fw inter.Framework, ctl *gomock.Controller, name string, ddlFile string) (*AgentHarness, error) {
    37  	ddl, err := addl.New(ddlFile)
    38  	if err != nil {
    39  		return nil, err
    40  	}
    41  
    42  	return New(fw, ctl, name, ddl)
    43  }
    44  
    45  // New creates a new test harness with all the actions found in ddl mocked using gomock
    46  func New(fw inter.Framework, ctl *gomock.Controller, name string, ddl *addl.DDL) (*AgentHarness, error) {
    47  	h := &AgentHarness{
    48  		name:    name,
    49  		ddl:     ddl,
    50  		fw:      fw,
    51  		log:     fw.Logger(fmt.Sprintf("%s_harnass", name)),
    52  		actions: make(map[string]*MockActionMiddleware),
    53  	}
    54  
    55  	if len(ddl.ActionNames()) == 0 {
    56  		return nil, fmt.Errorf("no actions defined")
    57  	}
    58  
    59  	for _, action := range ddl.Actions {
    60  		h.actions[action.Name] = NewMockActionMiddleware(ctl)
    61  	}
    62  
    63  	return h, nil
    64  }
    65  
    66  type ActionMiddleware interface {
    67  	Action(ctx context.Context, req *mcorpc.Request, rep *mcorpc.Reply, agent *mcorpc.Agent, conn inter.ConnectorInfo)
    68  }
    69  
    70  type AgentHarness struct {
    71  	name    string
    72  	ddl     *addl.DDL
    73  	fw      mcorpc.ChoriaFramework
    74  	log     *logrus.Entry
    75  	actions map[string]*MockActionMiddleware
    76  }
    77  
    78  // Stub sets an implementation action to be called for a specific action, panics for
    79  // unknown action
    80  func (h *AgentHarness) Stub(actionName string, action mcorpc.Action) *gomock.Call {
    81  	act, ok := h.actions[actionName]
    82  	if !ok {
    83  		panic(fmt.Sprintf("unknown action %s", actionName))
    84  	}
    85  
    86  	return act.EXPECT().Action(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Do(
    87  		func(ctx context.Context, req *mcorpc.Request, rep *mcorpc.Reply, agent *mcorpc.Agent, conn inter.ConnectorInfo) {
    88  			action(ctx, req, rep, agent, conn)
    89  		},
    90  	)
    91  }
    92  
    93  // Agent creates a valid mcorpc agent ready for passing to a server instance AgentManager
    94  func (h *AgentHarness) Agent() (*mcorpc.Agent, error) {
    95  	agent := mcorpc.New(h.name, h.ddl.Metadata, h.fw, h.log)
    96  
    97  	for n := range h.actions {
    98  		act := h.actions[n]
    99  		err := agent.RegisterAction(n, func(ctx context.Context, req *mcorpc.Request, rep *mcorpc.Reply, agent *mcorpc.Agent, conn inter.ConnectorInfo) {
   100  			defer ginkgo.GinkgoRecover()
   101  			act.Action(ctx, req, rep, agent, conn)
   102  		})
   103  		if err != nil {
   104  			return nil, err
   105  		}
   106  	}
   107  
   108  	return agent, nil
   109  }