[ondatra] Wait for gNOI server to be down after reboot request.

sdn_tests/pins_ondatra/bazel/patches/ondatra.patch

@@ -7,6 +7,19 @@ index 4d431d1..0ff43a4 100644
  	pathzpb "github.com/openconfig/gnsi/pathz"
 
 -	grpb "github.com/openconfig/gribi/v1/proto/service"
++	grpb "github.com/openconfig/gribi/proto/service"
+ 	opb "github.com/openconfig/ondatra/proto"
+ 	p4pb "github.com/p4lang/p4runtime/go/p4/v1"
+ )
+diff --git a/fakebind/fakebind.go b/fakebind/fakebind.go
+index 7dd0538..800a272 100644
+--- a/fakebind/fakebind.go
++++ b/fakebind/fakebind.go
+@@ -28,7 +28,7 @@ import (
+ 	"google.golang.org/grpc"
+
+ 	gpb "github.com/openconfig/gnmi/proto/gnmi"
+-	grpb "github.com/openconfig/gribi/v1/proto/service"
 +	grpb "github.com/openconfig/gribi/proto/service"
  	opb "github.com/openconfig/ondatra/proto"
  	p4pb "github.com/p4lang/p4runtime/go/p4/v1"
@@ -141,4 +154,4 @@ index 7145250..85cb489 100644
 +import "github.com/openconfig/ondatra/proto/testbed.proto";
 
  option go_package = "github.com/openconfig/ondatra/proto/reservation";
-
+

sdn_tests/pins_ondatra/infrastructure/testhelper/BUILD.bazel

@@ -1,5 +1,10 @@
 load("@io_bazel_rules_go//go:def.bzl", "go_library")
 
+def_visibility = [
+    "//platforms/networking/gpins/testing:__subpackages__",
+    "//third_party/pins_infra:__subpackages__",
+]
+
 package(
     default_visibility = ["//visibility:public"],
     licenses = ["notice"],
@@ -21,6 +26,7 @@ go_library(
         "platform_info.go",
         "port_management.go",
 	"results.go",
+        "utils.go",
         "ssh.go",
         "//infrastructure/testhelper/platform_info:platform_info",
     ],

sdn_tests/pins_ondatra/infrastructure/testhelper/config_convergence.go

@@ -174,26 +174,6 @@ func CompareConfigAndStateValues(ctx context.Context, t *testing.T, dut *ondatra
 	return r.String(), nil
 }
 
-// pollFunc returns true if the condition is met.
-type pollFunc func() bool
-
-// poll polls the condition until it is met or the context is done.
-func poll(ctx context.Context, t *testing.T, interval time.Duration, pf pollFunc) error {
-	ticker := time.NewTicker(interval)
-	defer ticker.Stop()
-	for {
-		select {
-		case <-ctx.Done():
-			return fmt.Errorf("polling for condition failed, err: %v", ctx.Err())
-		case <-ticker.C:
-			if pf() {
-				log.InfoContextf(ctx, "polling done")
-				return nil
-			}
-		}
-	}
-}
-
 // WaitForConfigConvergence checks for differences between config and state.
 // Polls till configConvergenceTimeout,
 // returns error if the difference still exists.
@@ -208,18 +188,18 @@ func WaitForConfigConvergence(ctx context.Context, t *testing.T, dut *ondatra.DU
 	defer cancel()
 	dutName := dut.Name()
 	// Poll until the config and state are similar.
-	return poll(ctx, t, configConvergencePollInterval, func() bool {
+	return poll(ctx, configConvergencePollInterval, func() pollStatus {
 		diff, err := CompareConfigAndStateValues(ctx, t, dut, config)
 		if err != nil {
 			log.InfoContextf(ctx, "Comparing config and state failed for dut: %v, err: %v", dutName, err)
-			return false
+			return continuePoll
 		}
 		if diff == "" {
 			log.InfoContextf(ctx, "Config and state converged for dut: %v", dutName)
-			return true
+			return exitPoll
 		}
 		log.InfoContextf(ctx, "diff in config and state found for dut: %v\ndiff_begin:\n%v\ndiff_end\n", dutName, diff)
-		return false
+		return continuePoll
 	})
 }
 
@@ -369,30 +349,30 @@ func WaitForSwitchState(ctx context.Context, t *testing.T, dut *ondatra.DUTDevic
 	defer cancel()
 
 	// Poll until the switch is ready or the context is done.
-	err := poll(ctx, t, switchStatePollInterval, func() bool {
+	err := poll(ctx, switchStatePollInterval, func() pollStatus {
 		switch s := switchState; s {
 		case down:
 			if err := GNOIAble(t, dut); err != nil {
 				log.InfoContextf(ctx, "GNOIAble(dut=%v) failed, err: %v", dutName, err)
-				return false
+				return continuePoll
 			}
 			switchState++
 		case gnoiAble:
 			if err := GNMIAble(t, dut); err != nil {
 				log.InfoContextf(ctx, "GNMIAble(dut=%v) failed, err: %v", dutName, err)
-				return false
+				return continuePoll
 			}
 			switchState++
 		case gnmiAble:
 			if err := WaitForAllPortsUp(ctx, t, dut); err != nil {
 				log.InfoContextf(ctx, "WaitForAllPortsUp(dut=%v) failed, err: %v", dutName, err)
-				return false
+				return continuePoll
 			}
 			switchState++
 		case portsUp:
-			return true
+			return exitPoll
 		}
-		return false
+		return continuePoll
 	})
 
 	if err == nil {

sdn_tests/pins_ondatra/infrastructure/testhelper/config_restorer.go

@@ -93,10 +93,16 @@ func NewConfigRestorerWithIgnorePaths(t *testing.T, ignorePaths []string) *Confi
 		t.Fatalf("config_restorer creation failed, errors: %v", err)
 	}
 
-        // Register a cleanup to restore the configs on test end.
-	t.Cleanup(func() {
-		cr.RestoreConfigsAndClose(t)
-	})
+        // Cleanup is only registered if running in a go test.
+	// As the cleanup callback would not be called otherwise.
+	if testing.Testing() { // True if running in a go test.
+		// Register a cleanup to restore the configs on test end.
+		log.InfoContextf(ctx, "Registering config_restorer for cleanup tasks.")
+		t.Cleanup(func() {
+			cr.RestoreConfigsAndClose(t)
+		})
+	}
+
 	return cr
 }
 
@@ -238,12 +244,10 @@ func (cr *ConfigRestorer) restoreConfigOnDiff(ctx context.Context, t *testing.T,
 
 // restoreReservedDevices tries to restore the config of the reserved devices
 // if the config differs from the saved config.
-func (cr *ConfigRestorer) restoreReservedDevices(t *testing.T) {
-        t.Helper()
+func (cr *ConfigRestorer) restoreReservedDevices(t *testing.T) error {
 	ctx := context.Background()
 	if cr.savedConfigs == nil {
-		log.InfoContextf(ctx, "configRestorer.savedConfigs is not initialized.")
-		return
+		return fmt.Errorf("savedConfigs are not initialized")
 	}
 
         wg := sync.WaitGroup{}
@@ -268,16 +272,24 @@ func (cr *ConfigRestorer) restoreReservedDevices(t *testing.T) {
 
         // Collect all the errors and fail the test on error.
 	if err := collectErrors(errCh); err != nil {
-		t.Fatalf("failed to restore config, errors: %v", err)
+		return fmt.Errorf("failed to restore config, errors: %v", err)
 	}
         log.InfoContextf(ctx, "Config restored for all the reserved devices.")
+        return nil
+}
+
+// RestoreConfigs restores the config of reserved devices.s
+func (cr *ConfigRestorer) RestoreConfigs(t *testing.T) error {
+	return cr.restoreReservedDevices(t)
 }
 
 // RestoreConfigsAndClose restores the config of reserved devices
 // and closes the configRestorer object.
 func (cr *ConfigRestorer) RestoreConfigsAndClose(t *testing.T) {
         t.Helper()
-	cr.restoreReservedDevices(t)
+	if err := cr.RestoreConfigs(t); err != nil {
+		t.Fatalf("config_restorer failed, errors: %v", err)
+	}
 	cr.savedConfigs = nil
 	cr.ignorePaths = nil
 }

sdn_tests/pins_ondatra/infrastructure/testhelper/gnoi.go

@@ -44,18 +44,20 @@ var (
 
 // RebootParams specify the reboot parameters used by the Reboot API.
 type RebootParams struct {
-	request       any
-	waitTime      time.Duration
-	checkInterval time.Duration
-	lmTTkrID      string // latency measurement testtracker UUID
-	lmTitle       string // latency measurement title
+	request        any
+	waitTime       time.Duration
+	checkInterval  time.Duration
+	requestTimeout time.Duration
+	lmTTkrID       string // latency measurement testtracker UUID
+	lmTitle        string // latency measurement title
 }
 
 // NewRebootParams returns RebootParams structure with default values.
 func NewRebootParams() *RebootParams {
 	return &RebootParams{
-		waitTime:      4 * time.Minute,
-		checkInterval: 20 * time.Second,
+		waitTime:       4 * time.Minute,
+		checkInterval:  20 * time.Second,
+		requestTimeout: 2 * time.Minute,
 	}
 }
 
@@ -82,6 +84,14 @@ func (p *RebootParams) WithRequest(r any) *RebootParams {
 	return p
 }
 
+// WithRequestTimeout adds the timeout for the reboot request.
+// The function will wait for the gNOI server to be down within this duration.
+// Default value is 2 minutes.
+func (p *RebootParams) WithRequestTimeout(timeout time.Duration) *RebootParams {
+	p.requestTimeout = timeout
+	return p
+}
+
 // WithLatencyMeasurement adds testtracker uuid and title for latency measurement.
 func (p *RebootParams) WithLatencyMeasurement(testTrackerID, title string) *RebootParams {
 	p.lmTTkrID = testTrackerID
@@ -132,28 +142,47 @@ func Reboot(t *testing.T, d *ondatra.DUTDevice, params *RebootParams) error {
 		return nil
 	}
 
-	log.Infof("Polling gNOI server reachability in %v intervals for max duration of %v", params.checkInterval, params.waitTime)
+	rebootRequestTimeout := params.requestTimeout
+	ctx, cancel := context.WithTimeout(context.Background(), rebootRequestTimeout)
+	defer cancel()
+
+	// The switch backend might not have processed the request or might take
+	// sometime to execute the request. So poll for the gNOI server to be down,
+	// or context to expire.
+	pollErr := poll(ctx, 10*time.Second /*(pollInterval)*/, func() pollStatus {
+		err := GNOIAble(t, d)
+		timeElapsed := (time.Now().UnixNano() - timeBeforeReboot) / int64(time.Second)
+		if err == nil {
+			log.Infof("gNOI server is still up after %v seconds", timeElapsed)
+			return continuePoll
+		}
+		log.Infof("gNOI server is down after %v seconds, got error while checking gNOI server reachability: %v as expected", timeElapsed, err)
+		return exitPoll
+	})
+	if pollErr != nil {
+		log.WarningContextf(ctx, "Polling gNOI server to be down within time: %v failed: %v", rebootRequestTimeout, pollErr)
+		log.InfoContextf(ctx, "Continue to check if the switch has rebooted")
+	}
+
+	log.InfoContextf(ctx, "Polling gNOI server reachability in %v intervals for max duration of %v", params.checkInterval, params.waitTime)
 	for timeout := time.Now().Add(params.waitTime); time.Now().Before(timeout); {
-		// The switch backend might not have processed the request or might take
-		// sometime to execute the request. So wait for check interval time and
-		// later verify that the switch rebooted within the specified wait time.
 		time.Sleep(params.checkInterval)
 		doneTime := time.Now()
 		timeElapsed := (doneTime.UnixNano() - timeBeforeReboot) / int64(time.Second)
 
 		if err := GNOIAble(t, d); err != nil {
-			log.Infof("gNOI server not up after %v seconds", timeElapsed)
+			log.InfoContextf(ctx, "gNOI server not up after %v seconds", timeElapsed)
 			continue
 		}
-		log.Infof("gNOI server up after %v seconds", timeElapsed)
+		log.InfoContextf(ctx, "gNOI server up after %v seconds", timeElapsed)
 
 		// An extra check to ensure that the system has rebooted.
 		if bootTime := gnmiSystemBootTimeGet(t, d); bootTime < uint64(timeBeforeReboot) {
-			log.Infof("Switch has not rebooted after %v seconds", timeElapsed)
+			log.InfoContextf(ctx, "Switch has not rebooted after %v seconds", timeElapsed)
 			continue
 		}
 
-		log.Infof("Switch rebooted after %v seconds", timeElapsed)
+		log.InfoContextf(ctx, "Switch rebooted after %v seconds", timeElapsed)
 		return nil
 	}
 

sdn_tests/pins_ondatra/infrastructure/testhelper/p4rt.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/openconfig/ondatra"
 	"github.com/openconfig/ondatra/gnmi"
+        "github.com/openconfig/ygnmi/ygnmi"
 	"github.com/pkg/errors"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/protobuf/encoding/prototext"
@@ -35,15 +36,29 @@ var (
 		return 0, errors.Errorf("failed to get port ID for port %v from switch", port)
 	}
 	testhelperDeviceIDGet = func(t *testing.T, d *ondatra.DUTDevice) (uint64, error) {
-		deviceInfo, present := gnmi.Lookup(t, d, gnmi.OC().Component(icName).IntegratedCircuit().State()).Val()
+		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
+		defer cancel()
+		yc, err := ygnmiClient(ctx, d)
+		if err != nil {
+			return 0, fmt.Errorf("failed to create ygnmi client, err: %v", err)
+		}
+		v, err := ygnmi.Lookup(ctx, yc, gnmi.OC().Component(icName).IntegratedCircuit().State())
+		if err != nil {
+			return 0, fmt.Errorf("failed to lookup device ID, err: %v", err)
+		}
+		deviceInfo, present := v.Val()
 		if present && deviceInfo.NodeId != nil {
 			return *deviceInfo.NodeId, nil
 		}
 		// Configure default device ID on the switch.
-		gnmi.Replace(t, d, gnmi.OC().Component(icName).IntegratedCircuit().NodeId().Config(), defaultDeviceID)
+		_, err = ygnmi.Replace(ctx, yc, gnmi.OC().Component(icName).IntegratedCircuit().NodeId().Config(), defaultDeviceID)
+		if err != nil {
+			return 0, fmt.Errorf("failed to configure default device ID, err: %v", err)
+		}
 		// Verify that default device ID has been configured and return that.
-		if got, want := gnmi.Get(t, d, gnmi.OC().Component(icName).IntegratedCircuit().NodeId().State()), defaultDeviceID; got != want {
-			return 0, errors.Errorf("failed to configure default device ID")
+		devID, err := ygnmi.Await(ctx, yc, gnmi.OC().Component(icName).IntegratedCircuit().NodeId().State(), defaultDeviceID, nil)
+		if err != nil {
+			return 0, fmt.Errorf("waiting for device ID to be %v failed, have %v, err: %v", defaultDeviceID, devID.String(), err)
 		}
 		return defaultDeviceID, nil
 	}

sdn_tests/pins_ondatra/infrastructure/testhelper/platform_components.go

@@ -14,7 +14,7 @@ import (
 
 // SwitchNameRegex returns the regex for switch name.
 func SwitchNameRegex() string {
-	return ""
+	return "^(ju|df|mn|ocs)(\\d+).*\\.([a-z]{3})(\\d{2})([a-z]?)\\.(net|prod).google.com$"
 }
 
 // ImageVersionRegex returns the regular expressions for the image version of the switch.