MINOR: Ignore unassigned records in MockConsumer (#21631)

In `poll()`, `MockConsumer` iterates the added `records`, and only
checks that they belong to an assigned partition quite late in the
process.

This can cause an error if records were added for a partition that was
later removed from the assignment:

```java
consumer.assign(Arrays.asList(new TopicPartition("t1", 0), new
TopicPartition("t2", 0)));
consumer.addRecord(new ConsumerRecord("t1", 0, 0, "a", 123));
consumer.addRecord(new ConsumerRecord("t2", 0, 0, "b", 123));

consumer.assign(Collections.singleton(new TopicPartition("t1", 0)));

// throws IllegalStateException
consumer.poll(Duration.seconds(1));
```

Moving this check earlier in the process avoids this error in `poll()`,
instead discarding records that do not belong to the assignment.

This enables tests to proactively add all their records and then modify
the assignment multiple times, if necessary, to test specific
combinations of partitions.

This is particularly useful in some of the `streams` tests, where we
need to test with multiple state store changelogs, but the assignment
gets changed internally by Kafka Streams (notably,
`GlobalStateManagerImplTest`).

Reviewers: Bill Bejeck <bbejeck@apache.org>

Author: nicktelford

clients/src/main/java/org/apache/kafka/clients/consumer/MockConsumer.java

@@ -284,7 +284,7 @@ public synchronized ConsumerRecords<K, V> poll(final Duration timeout) {
         while (partitionsIter.hasNext() && numPollRecords < this.maxPollRecords) {
             Map.Entry<TopicPartition, List<ConsumerRecord<K, V>>> entry = partitionsIter.next();
 
-            if (!subscriptions.isPaused(entry.getKey())) {
+            if (!subscriptions.isPaused(entry.getKey()) && subscriptions.isAssigned(entry.getKey())) {
                 final Iterator<ConsumerRecord<K, V>> recIterator = entry.getValue().iterator();
                 while (recIterator.hasNext()) {
                     if (numPollRecords >= this.maxPollRecords) {
@@ -298,7 +298,7 @@ public synchronized ConsumerRecords<K, V> poll(final Duration timeout) {
                         throw new OffsetOutOfRangeException(Collections.singletonMap(entry.getKey(), position));
                     }
 
-                    if (assignment().contains(entry.getKey()) && rec.offset() >= position) {
+                    if (rec.offset() >= position) {
                         results.computeIfAbsent(entry.getKey(), partition -> new ArrayList<>()).add(rec);
                         Metadata.LeaderAndEpoch leaderAndEpoch = new Metadata.LeaderAndEpoch(Optional.empty(), rec.leaderEpoch());
                         SubscriptionState.FetchPosition newPosition = new SubscriptionState.FetchPosition(

streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java

@@ -1456,7 +1456,7 @@ public void shouldNotReturnDataAfterTaskMigrated(final boolean processingThreads
 
         assertThat(thrown.getCause(), isA(IllegalStateException.class));
         // The Mock consumer shall throw as the assignment has been wiped out, but records are assigned.
-        assertEquals("No current assignment for partition topic1-1", thrown.getCause().getMessage());
+        assertEquals("Cannot add records for a partition that is not assigned to the consumer", thrown.getCause().getMessage());
         assertFalse(consumer.shouldRebalance());
 
         verify(taskManager).handleLostAll();