GODRIVER-2466 Implement assertEventCount

internal/integration/unified/client_entity.go

@@ -320,11 +320,31 @@ func (c *clientEntity) getEventCount(eventType monitoringEventType) int32 {
 	return c.eventsCount[eventType]
 }
 
-func (c *clientEntity) getServerDescriptionChangedEventCount(evt serverDescriptionChangedEventInfo) int32 {
+func (c *clientEntity) getServerDescriptionChangedEventCount(expected serverDescriptionChangedEventInfo) int32 {
 	c.serverDescriptionChangedEventsCountLock.Lock()
 	defer c.serverDescriptionChangedEventsCountLock.Unlock()
 
-	return c.serverDescriptionChangedEventsCount[evt]
+	// If the previousDescription or newDescription is "nil" in the expected case,
+	// then those fields can be anything.
+	for actual, count := range c.serverDescriptionChangedEventsCount {
+		actPrevD := actual.PreviousDescription
+		expPrevD := expected.PreviousDescription
+
+		if expPrevD != nil && expPrevD.Type != actPrevD.Type {
+			continue
+		}
+
+		actNewD := actual.NewDescription
+		expNewD := expected.NewDescription
+
+		if expNewD != nil && expNewD.Type != actNewD.Type {
+			continue
+		}
+
+		return count
+	}
+
+	return 0
 }
 
 func getSecondsSinceEpoch() float64 {

internal/integration/unified/event.go

@@ -11,7 +11,6 @@ import (
 
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/event"
-	"go.mongodb.org/mongo-driver/mongo/description"
 )
 
 type monitoringEventType string
@@ -124,21 +123,21 @@ type serverDescription struct {
 type serverDescriptionChangedEventInfo struct {
 	// NewDescription  corresponds to the server description as it was after
 	// the change that triggered this event.
-	NewDescription serverDescription
+	NewDescription *serverDescription
 
 	// PreviousDescription corresponds to the server description as it was
 	// before the change that triggered this event
-	PreviousDescription serverDescription
+	PreviousDescription *serverDescription
 }
 
 // newServerDescriptionChangedEventInfo returns a new serverDescriptionChangedEvent
 // instance for the given event.
 func newServerDescriptionChangedEventInfo(evt *event.ServerDescriptionChangedEvent) *serverDescriptionChangedEventInfo {
 	return &serverDescriptionChangedEventInfo{
-		NewDescription: serverDescription{
+		NewDescription: &serverDescription{
 			Type: evt.NewDescription.Kind.String(),
 		},
-		PreviousDescription: serverDescription{
+		PreviousDescription: &serverDescription{
 			Type: evt.PreviousDescription.Kind.String(),
 		},
 	}
@@ -156,21 +155,18 @@ func (evt *serverDescriptionChangedEventInfo) UnmarshalBSON(data []byte) error {
 		return err
 	}
 
-	// Set the default case to "Unknown" for the NewDescription.Type field.
-	evt.NewDescription.Type = description.TopologyKind(description.Unknown).String()
-
 	// Lookup the previous description, if any.
 	if newDescription, err := raw.LookupErr("newDescription"); err == nil {
-		evt.NewDescription.Type = newDescription.Document().Lookup("type").StringValue()
+		evt.NewDescription = &serverDescription{
+			Type: newDescription.Document().Lookup("type").StringValue(),
+		}
 	}
 
-	// Set the default case to "Unknown" for the PreviousDescription.Type
-	// field.
-	evt.PreviousDescription.Type = description.TopologyKind(description.Unknown).String()
-
 	// Lookup the previous description, if any.
 	if previousDescription, err := raw.LookupErr("previousDescription"); err == nil {
-		evt.PreviousDescription.Type = previousDescription.Document().Lookup("type").StringValue()
+		evt.PreviousDescription = &serverDescription{
+			Type: previousDescription.Document().Lookup("type").StringValue(),
+		}
 	}
 
 	return nil

internal/integration/unified/testrunner_operation.go

@@ -50,6 +50,45 @@ func (lp *loopArgs) iterationsStored() bool {
 	return lp.IterationsEntityID != ""
 }
 
+type assertEventCountArguments struct {
+	ClientID string              `bson:"client"`
+	Event    map[string]bson.Raw `bson:"event"`
+	Count    int32               `bson:"count"`
+}
+
+func assertEventCount(ctx context.Context, args assertEventCountArguments) error {
+	client, err := entities(ctx).client(args.ClientID)
+	if err != nil {
+		return err
+	}
+
+	for rawEventType, eventDoc := range args.Event {
+		eventType, ok := monitoringEventTypeFromString(rawEventType)
+		if !ok {
+			continue
+		}
+
+		var actualCount int32
+		switch eventType {
+		case serverDescriptionChangedEvent:
+			actualCount = getServerDescriptionChangedEventCount(client, eventDoc)
+			if actualCount == args.Count {
+				return nil
+			}
+		default:
+			actualCount = client.getEventCount(eventType)
+			if actualCount == args.Count {
+				return nil
+			}
+		}
+
+		return fmt.Errorf("expected event %q to have occured %v times, actual: %v",
+			rawEventType, args.Count, actualCount)
+	}
+
+	return nil
+}
+
 func executeTestRunnerOperation(ctx context.Context, op *operation, loopDone <-chan struct{}) error {
 	args := op.Arguments
 
@@ -145,6 +184,13 @@ func executeTestRunnerOperation(ctx context.Context, op *operation, loopDone <-c
 		coll := lookupString(args, "collectionName")
 		index := lookupString(args, "indexName")
 		return verifyIndexExists(ctx, db, coll, index, false)
+	case "assertEventCount":
+		var aecArgs assertEventCountArguments
+		if err := bson.Unmarshal(op.Arguments, &aecArgs); err != nil {
+			return fmt.Errorf("error unmarshalling event to assertEventCountArguments: %v", err)
+		}
+
+		return assertEventCount(ctx, aecArgs)
 	case "loop":
 		var unmarshaledArgs loopArgs
 		if err := bson.Unmarshal(args, &unmarshaledArgs); err != nil {