From 394752ba93fe132c2f45b62d37471c735aaa5e50 Mon Sep 17 00:00:00 2001
From: Karim Jahed <kjahed@cisco.com>
Date: Thu, 23 Jan 2025 14:13:48 -0500
Subject: [PATCH] use stream

---
 .../metadata.textproto                        |  2 +
 .../telemetry_basic_check_test.go             | 54 +++++++++++--------
 2 files changed, 35 insertions(+), 21 deletions(-)

diff --git a/feature/gnmi/otg_tests/telemetry_basic_check_test/metadata.textproto b/feature/gnmi/otg_tests/telemetry_basic_check_test/metadata.textproto
index 9918782cd47..ccae88e816d 100644
--- a/feature/gnmi/otg_tests/telemetry_basic_check_test/metadata.textproto
+++ b/feature/gnmi/otg_tests/telemetry_basic_check_test/metadata.textproto
@@ -13,6 +13,7 @@ platform_exceptions: {
   deviations: {
     ipv4_missing_enabled: true
     interface_counters_from_container: true
+    interface_counters_update_delayed: true
   }
 }
 platform_exceptions: {
@@ -24,6 +25,7 @@ platform_exceptions: {
     ipv4_missing_enabled: true
     interface_counters_from_container: true
     os_component_parent_is_supervisor_or_linecard: true
+    interface_counters_update_delayed: true
   }
 }
 platform_exceptions: {
diff --git a/feature/gnmi/otg_tests/telemetry_basic_check_test/telemetry_basic_check_test.go b/feature/gnmi/otg_tests/telemetry_basic_check_test/telemetry_basic_check_test.go
index 70286ac2d75..9b9c11582c8 100644
--- a/feature/gnmi/otg_tests/telemetry_basic_check_test/telemetry_basic_check_test.go
+++ b/feature/gnmi/otg_tests/telemetry_basic_check_test/telemetry_basic_check_test.go
@@ -768,37 +768,48 @@ func TestP4rtNodeID(t *testing.T) {
 	}
 }
 
-func fetchInAndOutPkts(t *testing.T, dut *ondatra.DUTDevice, dp1, dp2 *ondatra.Port) (uint64, uint64) {
+func fetchInAndOutPkts(t *testing.T, dut *ondatra.DUTDevice, dp1, dp2 *ondatra.Port,
+	inTarget, outTarget uint64) (uint64, uint64) {
 	t.Helper()
-	inPktStream := samplestream.New(t, dut, gnmi.OC().Interface(dp1.Name()).Counters().InUnicastPkts().State(), 60*time.Second)
+
+	inPktStream := samplestream.New(t, dut, gnmi.OC().Interface(dp1.Name()).Counters().InUnicastPkts().State(), 10*time.Second)
 	defer inPktStream.Close()
-	outPktStream := samplestream.New(t, dut, gnmi.OC().Interface(dp2.Name()).Counters().OutUnicastPkts().State(), 60*time.Second)
+	outPktStream := samplestream.New(t, dut, gnmi.OC().Interface(dp2.Name()).Counters().OutUnicastPkts().State(), 10*time.Second)
 	defer outPktStream.Close()
 
 	var wg sync.WaitGroup
 	var inPktsV, outPktsV uint64
 
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		if v := inPktStream.Next(); v != nil {
+	startTime := time.Now()
+	timeout := 10 * time.Second
+	if deviations.InterfaceCountersUpdateDelayed(dut) {
+		timeout = 30 * time.Second
+	}
+
+	for {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if v := inPktStream.Next(); v != nil {
+				if val, ok := v.Val(); ok {
+					inPktsV = val
+				}
+			}
+		}()
+		if v := outPktStream.Next(); v != nil {
 			if val, ok := v.Val(); ok {
-				inPktsV = val
+				outPktsV = val
 			}
 		}
-	}()
-	if v := outPktStream.Next(); v != nil {
-		if val, ok := v.Val(); ok {
-			outPktsV = val
+		wg.Wait()
+
+		if inPktsV >= inTarget && outPktsV >= outTarget {
+			break
 		}
-	}
-	wg.Wait()
 
-	if inPktsV == 0 {
-		t.Fatalf("Did not receive a value for in packet counter")
-	}
-	if outPktsV == 0 {
-		t.Fatalf("Did not receive a value for out packet counter")
+		if time.Since(startTime) > timeout {
+			t.Fatalf("Did not receive a packet counters in time")
+		}
 	}
 
 	return inPktsV, outPktsV
@@ -852,7 +863,7 @@ func TestIntfCounterUpdate(t *testing.T) {
 	otgutils.WaitForARP(t, ate.OTG(), config, "IPv4")
 
 	t.Log("Running traffic on DUT interfaces: ", dp1, dp2)
-	dutInPktsBeforeTraffic, dutOutPktsBeforeTraffic := fetchInAndOutPkts(t, dut, dp1, dp2)
+	dutInPktsBeforeTraffic, dutOutPktsBeforeTraffic := fetchInAndOutPkts(t, dut, dp1, dp2, 0, 0)
 	t.Log("inPkts and outPkts counters before traffic: ", dutInPktsBeforeTraffic, dutOutPktsBeforeTraffic)
 	otg.StartTraffic(t)
 	time.Sleep(10 * time.Second)
@@ -885,7 +896,8 @@ func TestIntfCounterUpdate(t *testing.T) {
 	if lossPct >= 0.1 {
 		t.Errorf("Get(traffic loss for flow %q: got %v, want < 0.1", flowName, lossPct)
 	}
-	dutInPktsAfterTraffic, dutOutPktsAfterTraffic := fetchInAndOutPkts(t, dut, dp1, dp2)
+	dutInPktsAfterTraffic, dutOutPktsAfterTraffic := fetchInAndOutPkts(t, dut, dp1, dp2,
+		dutInPktsBeforeTraffic+ateOutPkts, dutOutPktsBeforeTraffic+ateInPkts)
 	t.Log("inPkts and outPkts counters after traffic: ", dutInPktsAfterTraffic, dutOutPktsAfterTraffic)
 
 	if got, want := dutInPktsAfterTraffic-dutInPktsBeforeTraffic, ateOutPkts; got < want {