github.com/aporeto-inc/trireme-lib@v10.358.0+incompatible/controller/internal/enforcer/nfqdatapath/test_utils.go (about)

     1  package nfqdatapath
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  
     7  	"github.com/golang/mock/gomock"
     8  	"go.aporeto.io/enforcerd/trireme-lib/collector"
     9  	"go.aporeto.io/enforcerd/trireme-lib/controller/pkg/secrets/mocksecrets"
    10  )
    11  
    12  type endpointTypeMatcher struct {
    13  	x           interface{}
    14  	baseMatcher *myMatcher
    15  }
    16  
    17  func (m *endpointTypeMatcher) Matches(x interface{}) bool {
    18  	f1 := m.x.(*collector.FlowRecord)
    19  	f2 := x.(*collector.FlowRecord)
    20  
    21  	defaultChecks := f1.Destination.Type == f2.Destination.Type &&
    22  		f1.Destination.ID == f2.Destination.ID &&
    23  		f1.Source.Type == f2.Source.Type &&
    24  		f1.Source.ID == f2.Source.ID
    25  
    26  	if m.baseMatcher != nil {
    27  		return defaultChecks && m.baseMatcher.Matches(x)
    28  	}
    29  
    30  	return defaultChecks
    31  }
    32  
    33  func (m *endpointTypeMatcher) String() string {
    34  
    35  	out, err := json.Marshal(m.x)
    36  	if err != nil {
    37  		panic(err)
    38  	}
    39  
    40  	return fmt.Sprintf("is equal to %v", string(out))
    41  }
    42  
    43  // EndpointTypeMatcher extends MyMatcher to match endpoint Type and ID
    44  func EndpointTypeMatcher(x interface{}) gomock.Matcher {
    45  	return gomock.GotFormatterAdapter(&myGotFormatter{}, &endpointTypeMatcher{x: x, baseMatcher: &myMatcher{x: x}})
    46  }
    47  
    48  type myMatcher struct {
    49  	x interface{}
    50  }
    51  
    52  func (m *myMatcher) Matches(x interface{}) bool {
    53  	f1 := m.x.(*collector.FlowRecord)
    54  	f2 := x.(*collector.FlowRecord)
    55  
    56  	defaultChecks := f1.Destination.IP == f2.Destination.IP &&
    57  		f1.Source.IP == f2.Source.IP &&
    58  		f1.Destination.Port == f2.Destination.Port &&
    59  		f1.Action == f2.Action &&
    60  		f1.Count == f2.Count &&
    61  		f1.DropReason == f2.DropReason
    62  
    63  	return defaultChecks
    64  }
    65  
    66  func (m *myMatcher) String() string {
    67  
    68  	f := m.x.(*collector.FlowRecord)
    69  	return fmt.Sprintf("%d, %v, %v, %d, %d, %s", f.Count, f.Source.IP, f.Destination.IP, f.Destination.Port, f.Action, f.DropReason)
    70  }
    71  
    72  type myGotFormatter struct{}
    73  
    74  func (g *myGotFormatter) Got(got interface{}) string {
    75  
    76  	f := got.(*collector.FlowRecord)
    77  	return fmt.Sprintf("%d, %v, %v, %d, %d, %s", f.Count, f.Source.IP, f.Destination.IP, f.Destination.Port, f.Action, f.DropReason)
    78  }
    79  
    80  // MyMatcher returns gomock matcher
    81  func MyMatcher(x interface{}) gomock.Matcher {
    82  	return gomock.GotFormatterAdapter(&myGotFormatter{}, &myMatcher{x: x})
    83  }
    84  
    85  type packetEventMatcher struct {
    86  	x interface{}
    87  }
    88  
    89  func (p *packetEventMatcher) Matches(x interface{}) bool {
    90  	f1 := p.x.(*collector.PacketReport)
    91  	f2 := x.(*collector.PacketReport)
    92  	return f1.DestinationIP == f2.DestinationIP
    93  }
    94  
    95  func (p *packetEventMatcher) String() string {
    96  	return fmt.Sprintf("is equal to %v", p.x)
    97  }
    98  
    99  // PacketEventMatcher return gomock matcher
   100  func PacketEventMatcher(x interface{}) gomock.Matcher {
   101  	return &packetEventMatcher{x: x}
   102  }
   103  
   104  type myCounterMatcher struct {
   105  	x *collector.CounterReport
   106  }
   107  
   108  func (m *myCounterMatcher) Matches(x interface{}) bool {
   109  
   110  	f := x.(*collector.CounterReport)
   111  	if f.Namespace != "/ns1" {
   112  		return true
   113  	}
   114  	return m.x.PUID == f.PUID && m.x.Namespace == f.Namespace
   115  }
   116  
   117  func (m *myCounterMatcher) String() string {
   118  	return fmt.Sprintf("is equal to %v", m.x)
   119  }
   120  
   121  // MyCounterMatcher custom matcher for counter record
   122  func MyCounterMatcher(x *collector.CounterReport) gomock.Matcher {
   123  	return &myCounterMatcher{x: x}
   124  }
   125  
   126  type fakeSecrets struct {
   127  	id string
   128  	*mocksecrets.MockSecrets
   129  }
   130  
   131  func (f *fakeSecrets) setID(id string) {
   132  	f.id = id
   133  }
   134  
   135  func (f *fakeSecrets) getID() string {
   136  	return f.id
   137  }