fix(pubsub/pstest): Clear Subscription when calling ClearMessages.

mikeklaas · hongalex · commit 6de8edada13c · 2023-05-02T13:42:36.000-07:00
Currently, `ClearMessages()` only clears `Server` messages, not touching any messages that have been published to a Subscription but not yet delivered.  This can cause odd behaviour if those messages are delivered as `pstest` assumes that any such messages exist in the Server's `msgsByID` map.  Attempting to ModAck such a message results in a NPE.
diff --git a/pubsub/pstest/fake.go b/pubsub/pstest/fake.go
@@ -288,6 +288,9 @@ func (s *Server) ClearMessages() {
 	s.GServer.mu.Lock()
 	s.GServer.msgs = nil
 	s.GServer.msgsByID = make(map[string]*Message)
+	for _, sub := range s.GServer.subs {
+		sub.msgs = map[string]*message{}
+	}
 	s.GServer.mu.Unlock()
 }
 
diff --git a/pubsub/pstest/fake_test.go b/pubsub/pstest/fake_test.go
@@ -472,13 +472,19 @@ func TestPublishOrdered(t *testing.T) {
 }
 
 func TestClearMessages(t *testing.T) {
-	s := NewServer()
-	defer s.Close()
+	pclient, sclient, s, cleanup := newFake(context.TODO(), t)
+	defer cleanup()
+
+	top := mustCreateTopic(context.TODO(), t, pclient, &pb.Topic{Name: "projects/P/topics/T"})
+	sub := mustCreateSubscription(context.TODO(), t, sclient, &pb.Subscription{
+		Name:               "projects/P/subscriptions/S",
+		Topic:              top.Name,
+		AckDeadlineSeconds: 10,
+	})
 
 	for i := 0; i < 3; i++ {
-		s.Publish("projects/p/topics/t", []byte("hello"), nil)
+		s.Publish(top.Name, []byte("hello"), nil)
 	}
-	s.Wait()
 	msgs := s.Messages()
 	if got, want := len(msgs), 3; got != want {
 		t.Errorf("got %d messages, want %d", got, want)
@@ -488,6 +494,14 @@ func TestClearMessages(t *testing.T) {
 	if got, want := len(msgs), 0; got != want {
 		t.Errorf("got %d messages, want %d", got, want)
 	}
+
+	res, err := sclient.Pull(context.Background(), &pb.PullRequest{Subscription: sub.Name})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(res.ReceivedMessages) != 0 {
+		t.Errorf("got %d messages, want zero", len(res.ReceivedMessages))
+	}
 }
 
 // Note: this sets the fake's "now" time, so it is sensitive to concurrent changes to "now".

Original file line number	Diff line number	Diff line change
`@@ -288,6 +288,9 @@ func (s *Server) ClearMessages() {`
`288`	`288`	`s.GServer.mu.Lock()`
`289`	`289`	`s.GServer.msgs = nil`
`290`	`290`	`s.GServer.msgsByID = make(map[string]*Message)`
	`291`	`+ for _, sub := range s.GServer.subs {`
	`292`	`+ sub.msgs = map[string]*message{}`
	`293`	`+ }`
`291`	`294`	`s.GServer.mu.Unlock()`
`292`	`295`	`}`
`293`	`296`