HIVE-29842: fixing flaky TestYarnQueueMetricsCollector tests

ql/src/test/org/apache/hadoop/hive/ql/exec/tez/monitoring/yarnqueue/TestYarnQueueMetricsCollector.java

@@ -107,11 +107,24 @@ private void waitForInvocationCount(Object mock, int minCount, long timeoutMs) {
     }
   }
 
+  /**
+   * Awaits the first refresh poll and asserts it happened. Use before shutting down a
+   * collector whose per-test stubs are consumed only by the refresh thread.
+   */
+  private void awaitFirstPoll() {
+    waitForInvocationCount(mockYarnClient, 1, WAIT_TIMEOUT_MS);
+    assertTrue("Refresh task should have polled YARN",
+        mockingDetails(mockYarnClient).getInvocations().size() >= 1);
+  }
+
   /**
    * Configures mock objects with standard happy-path values.
    * Called from {@code @Before} so all tests start with a consistent baseline.
    * Stubs are lenient so tests that don't exercise these mocks don't fail with
    * UnnecessaryStubbingException.
+   * Conversely, a per-test (strict) stub consumed only by the async refresh thread must be
+   * awaited ({@code waitForInvocationCount}/{@code waitForSnapshot}) before the collector is
+   * shut down, or the stub can go unused and strict-stub validation fails the whole class.
    */
   private void setupHappyPathMocks() throws Exception {
     lenient().when(mockQueueStats.getAllocatedMemoryMB()).thenReturn(1024L);
@@ -169,6 +182,10 @@ public void testMetricsCollectionWithNullQueueInfo() throws Exception {
 
     YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "nonexistent", 10000, "test-query-2");
     try {
+      // The pre-first-poll state (also a null snapshot) is not deterministically observable,
+      // as the refresh task runs with no initial delay; await the poll so the null asserted
+      // below is attributable to the stubbed null QueueInfo, not to "no poll yet".
+      awaitFirstPoll();
       assertNull("Snapshot should be null for nonexistent queue", collector.getLatestSnapshot());
     } finally {
       collector.shutdown();
@@ -221,9 +238,7 @@ public void testPercentageCalculationWithZeroTotal() {
   }
 
   @Test
-  public void testShutdownIdempotency() throws Exception {
-    when(mockYarnClient.getQueueInfo("default")).thenReturn(mockQueueInfo);
-
+  public void testShutdownIdempotency() {
     YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "default", 10000, "test-query-4");
     collector.shutdown();
     collector.shutdown(); // second call must be safe
@@ -237,18 +252,15 @@ public void testExceptionDuringCollection() throws Exception {
 
     YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "default", 10000, "test-query-5");
     try {
+      awaitFirstPoll();
       assertNull("Snapshot should be null after exception", collector.getLatestSnapshot());
     } finally {
       collector.shutdown();
     }
   }
 
   @Test
-  public void testQueueNameRetrieval() throws Exception {
-    when(mockYarnClient.getQueueInfo(anyString())).thenReturn(mockQueueInfo);
-    when(mockQueueInfo.getQueueStatistics()).thenReturn(null);
-    when(mockQueueInfo.getCapacity()).thenReturn(0.5f);
-
+  public void testQueueNameRetrieval() {
     YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "production", 10000, "test-query-6");
     try {
       assertEquals("Queue name should match", "production", collector.getQueueName());
@@ -305,6 +317,7 @@ public void testExecutorCleanupOnInitializationFailure() throws Exception {
 
     YarnQueueMetricsCollector collector = newCollector(mockYarnClient, "default", 10000, "init-fail-query");
     try {
+      awaitFirstPoll();
       assertNull("Snapshot should be null after init failure", collector.getLatestSnapshot());
     } finally {
       collector.shutdown();