diff --git a/tests/integration/system/nsf/BUILD.bazel b/tests/integration/system/nsf/BUILD.bazel
index 4b094855..b539b114 100644
--- a/tests/integration/system/nsf/BUILD.bazel
+++ b/tests/integration/system/nsf/BUILD.bazel
@@ -31,6 +31,7 @@ cc_library(
         "//lib/utils:generic_testbed_utils",
         "//lib/validator:validator_lib",
         "//tests/integration/system/nsf/interfaces:component_validator",
+        "//tests/integration/system/nsf/interfaces:testbed",
         "//thinkit:generic_testbed",
         "//thinkit:generic_testbed_fixture",
         "//thinkit:mirror_testbed",
@@ -61,8 +62,8 @@ cc_library(
         "//tests/integration/system/nsf/interfaces:component_validator",
         "//tests/integration/system/nsf/interfaces:flow_programmer",
         "//tests/integration/system/nsf/interfaces:test_params",
+        "//tests/integration/system/nsf/interfaces:testbed",
         "//tests/integration/system/nsf/interfaces:traffic_helper",
-        "//thinkit:generic_testbed",
         "//thinkit:ssh_client",
         "//thinkit/proto:generic_testbed_cc_proto",
         "@com_github_google_glog//:glog",
diff --git a/tests/integration/system/nsf/component_validators/BUILD.bazel b/tests/integration/system/nsf/component_validators/BUILD.bazel
index 8ae06495..b6394cac 100644
--- a/tests/integration/system/nsf/component_validators/BUILD.bazel
+++ b/tests/integration/system/nsf/component_validators/BUILD.bazel
@@ -29,7 +29,7 @@ cc_library(
     hdrs = ["sai_validator.h"],
     deps = [
         "//tests/integration/system/nsf/interfaces:component_validator",
-        "//thinkit:generic_testbed",
+        "//tests/integration/system/nsf/interfaces:testbed",
         "@com_github_google_glog//:glog",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings:string_view",
@@ -42,7 +42,7 @@ cc_library(
     hdrs = ["swss_validator.h"],
     deps = [
         "//tests/integration/system/nsf/interfaces:component_validator",
-        "//thinkit:generic_testbed",
+        "//tests/integration/system/nsf/interfaces:testbed",
         "@com_github_google_glog//:glog",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings:string_view",
@@ -55,7 +55,7 @@ cc_library(
     hdrs = ["syncd_validator.h"],
     deps = [
         "//tests/integration/system/nsf/interfaces:component_validator",
-        "//thinkit:generic_testbed",
+        "//tests/integration/system/nsf/interfaces:testbed",
         "@com_github_google_glog//:glog",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings:string_view",
diff --git a/tests/integration/system/nsf/component_validators/sai_validator.h b/tests/integration/system/nsf/component_validators/sai_validator.h
index 61071d4f..7d625ec3 100644
--- a/tests/integration/system/nsf/component_validators/sai_validator.h
+++ b/tests/integration/system/nsf/component_validators/sai_validator.h
@@ -19,7 +19,7 @@
 #include "absl/strings/string_view.h"
 #include "glog/logging.h"
 #include "tests/integration/system/nsf/interfaces/component_validator.h"
-#include "thinkit/generic_testbed.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 
 namespace pins_test {
 
@@ -27,13 +27,12 @@ namespace pins_test {
 // Component owners need to have their own implementations and register to be
 // used by the NSF Upgrade tests.
 class SaiValidator : public ComponentValidator {
-  absl::Status OnInit(absl::string_view version,
-                      thinkit::GenericTestbed& testbed) override {
+  absl::Status OnInit(absl::string_view version, Testbed& testbed) override {
     LOG(INFO) << "Sai Init";
     return absl::OkStatus();
   }
   absl::Status OnFlowCleanup(absl::string_view version,
-                             thinkit::GenericTestbed& testbed) override {
+                             Testbed& testbed) override {
     LOG(INFO) << "Sai Flow Cleanup";
     return absl::OkStatus();
   }
diff --git a/tests/integration/system/nsf/component_validators/swss_validator.h b/tests/integration/system/nsf/component_validators/swss_validator.h
index 3a9c46b3..dbf90e39 100644
--- a/tests/integration/system/nsf/component_validators/swss_validator.h
+++ b/tests/integration/system/nsf/component_validators/swss_validator.h
@@ -19,7 +19,7 @@
 #include "absl/strings/string_view.h"
 #include "glog/logging.h"
 #include "tests/integration/system/nsf/interfaces/component_validator.h"
-#include "thinkit/generic_testbed.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 
 namespace pins_test {
 
@@ -27,13 +27,12 @@ namespace pins_test {
 // Component owners need to have their own implementations and register to be
 // used by the NSF Upgrade tests.
 class SwssValidator : public ComponentValidator {
-  absl::Status OnInit(absl::string_view version,
-                      thinkit::GenericTestbed& testbed) override {
+  absl::Status OnInit(absl::string_view version, Testbed& testbed) override {
     LOG(INFO) << "Swss Init";
     return absl::OkStatus();
   }
   absl::Status OnFlowProgram(absl::string_view version,
-                             thinkit::GenericTestbed& testbed) override {
+                             Testbed& testbed) override {
     LOG(INFO) << "Swss Flow Program";
     return absl::OkStatus();
   }
diff --git a/tests/integration/system/nsf/component_validators/syncd_validator.h b/tests/integration/system/nsf/component_validators/syncd_validator.h
index 73a4d9c1..cef6a918 100644
--- a/tests/integration/system/nsf/component_validators/syncd_validator.h
+++ b/tests/integration/system/nsf/component_validators/syncd_validator.h
@@ -19,7 +19,7 @@
 #include "absl/strings/string_view.h"
 #include "glog/logging.h"
 #include "tests/integration/system/nsf/interfaces/component_validator.h"
-#include "thinkit/generic_testbed.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 
 namespace pins_test {
 
@@ -27,18 +27,17 @@ namespace pins_test {
 // Component owners need to have their own implementations and register to be
 // used by the NSF Upgrade tests.
 class SyncdValidator : public ComponentValidator {
-  absl::Status OnInit(absl::string_view version,
-                      thinkit::GenericTestbed& testbed) override {
+  absl::Status OnInit(absl::string_view version, Testbed& testbed) override {
     LOG(INFO) << "Syncd Init";
     return absl::OkStatus();
   }
   absl::Status OnFlowProgram(absl::string_view version,
-                             thinkit::GenericTestbed& testbed) override {
+                             Testbed& testbed) override {
     LOG(INFO) << "Syncd Flow Program";
     return absl::OkStatus();
   }
   absl::Status OnFlowCleanup(absl::string_view version,
-                             thinkit::GenericTestbed& testbed) override {
+                             Testbed& testbed) override {
     LOG(INFO) << "Syncd Flow Cleanup";
     return absl::OkStatus();
   }
diff --git a/tests/integration/system/nsf/flow_programmers/BUILD.bazel b/tests/integration/system/nsf/flow_programmers/BUILD.bazel
index 1f3344aa..d627a9ea 100644
--- a/tests/integration/system/nsf/flow_programmers/BUILD.bazel
+++ b/tests/integration/system/nsf/flow_programmers/BUILD.bazel
@@ -25,7 +25,7 @@ cc_library(
     hdrs = ["static_programmer.h"],
     deps = [
         "//tests/integration/system/nsf/interfaces:flow_programmer",
-        "//thinkit:generic_testbed",
+        "//tests/integration/system/nsf/interfaces:testbed",
         "@com_google_absl//absl/status",
     ],
 )
@@ -36,7 +36,7 @@ cc_library(
     hdrs = ["fuzzer_replay_programmer.h"],
     deps = [
         "//tests/integration/system/nsf/interfaces:flow_programmer",
-        "//thinkit:generic_testbed",
+        "//tests/integration/system/nsf/interfaces:testbed",
         "@com_google_absl//absl/status",
     ],
 )
diff --git a/tests/integration/system/nsf/flow_programmers/fuzzer_replay_programmer.h b/tests/integration/system/nsf/flow_programmers/fuzzer_replay_programmer.h
index 419f5801..ba79cd8a 100644
--- a/tests/integration/system/nsf/flow_programmers/fuzzer_replay_programmer.h
+++ b/tests/integration/system/nsf/flow_programmers/fuzzer_replay_programmer.h
@@ -17,16 +17,16 @@
 
 #include "absl/status/status.h"
 #include "tests/integration/system/nsf/interfaces/flow_programmer.h"
-#include "thinkit/generic_testbed.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 
 namespace pins_test {
 
 class FuzzerReplayProgrammer : public FlowProgrammer {
  public:
-  absl::Status ProgramFlows(thinkit::GenericTestbed& testbed) override {
+  absl::Status ProgramFlows(Testbed& testbed) override {
     return absl::OkStatus();
   };
-  absl::Status ClearFlows(thinkit::GenericTestbed& testbed) override {
+  absl::Status ClearFlows(Testbed& testbed) override {
     return absl::OkStatus();
   };
 };
diff --git a/tests/integration/system/nsf/flow_programmers/static_programmer.h b/tests/integration/system/nsf/flow_programmers/static_programmer.h
index 87b5d626..4aa3434a 100644
--- a/tests/integration/system/nsf/flow_programmers/static_programmer.h
+++ b/tests/integration/system/nsf/flow_programmers/static_programmer.h
@@ -17,16 +17,16 @@
 
 #include "absl/status/status.h"
 #include "tests/integration/system/nsf/interfaces/flow_programmer.h"
-#include "thinkit/generic_testbed.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 
 namespace pins_test {
 
 class StaticProgrammer : public FlowProgrammer {
  public:
-  absl::Status ProgramFlows(thinkit::GenericTestbed& testbed) override {
+  absl::Status ProgramFlows(Testbed& testbed) override {
     return absl::OkStatus();
   };
-  absl::Status ClearFlows(thinkit::GenericTestbed& testbed) override {
+  absl::Status ClearFlows(Testbed& testbed) override {
     return absl::OkStatus();
   };
 };
diff --git a/tests/integration/system/nsf/interfaces/BUILD.bazel b/tests/integration/system/nsf/interfaces/BUILD.bazel
index e123e64b..4cddee04 100644
--- a/tests/integration/system/nsf/interfaces/BUILD.bazel
+++ b/tests/integration/system/nsf/interfaces/BUILD.bazel
@@ -19,21 +19,34 @@ package(
     licenses = ["notice"],
 )
 
+cc_library(
+    name = "testbed",
+    testonly = True,
+    hdrs = ["testbed.h"],
+    deps = [
+        "//thinkit:generic_testbed",
+        "//thinkit:generic_testbed_fixture",
+        "//thinkit:mirror_testbed",
+        "//thinkit:mirror_testbed_fixture",
+    ],
+)
+
 cc_library(
     name = "flow_programmer",
     testonly = True,
     hdrs = ["flow_programmer.h"],
     deps = [
-        "//thinkit:generic_testbed",
+        ":testbed",
         "@com_google_absl//absl/status",
     ],
 )
 
 cc_library(
     name = "traffic_helper",
+    testonly = True,
     hdrs = ["traffic_helper.h"],
     deps = [
-        "//thinkit:generic_testbed",
+        ":testbed",
         "@com_google_absl//absl/status",
     ],
 )
@@ -43,7 +56,7 @@ cc_library(
     testonly = True,
     hdrs = ["component_validator.h"],
     deps = [
-        "//thinkit:generic_testbed",
+        ":testbed",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings:string_view",
     ],
@@ -56,8 +69,8 @@ cc_library(
     deps = [
         ":component_validator",
         ":flow_programmer",
+        ":testbed",
         ":traffic_helper",
-        "//tests/integration/system/nsf:util",
         "//thinkit:ssh_client",
     ],
 )
diff --git a/tests/integration/system/nsf/interfaces/component_validator.h b/tests/integration/system/nsf/interfaces/component_validator.h
index f92235f7..68e107b7 100644
--- a/tests/integration/system/nsf/interfaces/component_validator.h
+++ b/tests/integration/system/nsf/interfaces/component_validator.h
@@ -17,7 +17,7 @@
 
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
-#include "thinkit/generic_testbed.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 
 namespace pins_test {
 
@@ -50,52 +50,50 @@ class ComponentValidator {
   virtual ~ComponentValidator() = default;
 
   // Called before starting every NSF test.
-  virtual absl::Status OnInit(absl::string_view version,
-                              thinkit::GenericTestbed& testbed) {
+  virtual absl::Status OnInit(absl::string_view version, Testbed& testbed) {
     return absl::OkStatus();
   }
 
   // Called after programming flows on SUT.
   virtual absl::Status OnFlowProgram(absl::string_view version,
-                                     thinkit::GenericTestbed& testbed) {
+                                     Testbed& testbed) {
     return absl::OkStatus();
   }
 
   // Called after starting traffic from the Control Device or the Traffic
   // Generator in the testbed.
   virtual absl::Status OnStartTraffic(absl::string_view version,
-                                      thinkit::GenericTestbed& testbed) {
+                                      Testbed& testbed) {
     return absl::OkStatus();
   }
 
   // Called after an upgrade is performed on the SUT.
-  virtual absl::Status OnUpgrade(absl::string_view version,
-                                 thinkit::GenericTestbed& testbed) {
+  virtual absl::Status OnUpgrade(absl::string_view version, Testbed& testbed) {
     return absl::OkStatus();
   }
 
   // Called after a successful NSF reboot of the SUT.
   virtual absl::Status OnNsfReboot(absl::string_view version,
-                                   thinkit::GenericTestbed& testbed) {
+                                   Testbed& testbed) {
     return absl::OkStatus();
   }
 
   // Called after pushing config on the SUT.
   virtual absl::Status OnConfigPush(absl::string_view version,
-                                    thinkit::GenericTestbed& testbed) {
+                                    Testbed& testbed) {
     return absl::OkStatus();
   }
 
   // Called after stopping traffic from the Control Device or the Traffic
   // Generator in the testbed.
   virtual absl::Status OnStopTraffic(absl::string_view version,
-                                     thinkit::GenericTestbed& testbed) {
+                                     Testbed& testbed) {
     return absl::OkStatus();
   }
 
   // Called after clearing up flows from the SUT.
   virtual absl::Status OnFlowCleanup(absl::string_view version,
-                                     thinkit::GenericTestbed& testbed) {
+                                     Testbed& testbed) {
     return absl::OkStatus();
   }
 };
diff --git a/tests/integration/system/nsf/interfaces/flow_programmer.h b/tests/integration/system/nsf/interfaces/flow_programmer.h
index e87b0783..4ec1d24f 100644
--- a/tests/integration/system/nsf/interfaces/flow_programmer.h
+++ b/tests/integration/system/nsf/interfaces/flow_programmer.h
@@ -16,7 +16,7 @@
 #define PINS_TESTS_INTEGRATION_SYSTEM_NSF_INTERFACES_FLOW_PROGRAMMER_H_
 
 #include "absl/status/status.h"
-#include "thinkit/generic_testbed.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 
 namespace pins_test {
 
@@ -26,11 +26,12 @@ class FlowProgrammer {
  public:
   virtual ~FlowProgrammer() = default;
 
-  // Programs a predefined flow on the SUT.
-  virtual absl::Status ProgramFlows(thinkit::GenericTestbed& testbed) = 0;
+  // Programs a predefined flow on the SUT based on the given IP version and
+  // protocol.
+  virtual absl::Status ProgramFlows(Testbed& testbed) = 0;
 
   // Clears all flows on the SUT.
-  virtual absl::Status ClearFlows(thinkit::GenericTestbed& testbed) = 0;
+  virtual absl::Status ClearFlows(Testbed& testbed) = 0;
 };
 
 }  // namespace pins_test
diff --git a/tests/integration/system/nsf/interfaces/test_params.h b/tests/integration/system/nsf/interfaces/test_params.h
index bb50287b..ff1fa503 100644
--- a/tests/integration/system/nsf/interfaces/test_params.h
+++ b/tests/integration/system/nsf/interfaces/test_params.h
@@ -22,8 +22,8 @@
 
 #include "tests/integration/system/nsf/interfaces/component_validator.h"
 #include "tests/integration/system/nsf/interfaces/flow_programmer.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 #include "tests/integration/system/nsf/interfaces/traffic_helper.h"
-#include "tests/integration/system/nsf/util.h"
 #include "thinkit/ssh_client.h"
 
 namespace pins_test {
diff --git a/tests/integration/system/nsf/interfaces/testbed.h b/tests/integration/system/nsf/interfaces/testbed.h
new file mode 100644
index 00000000..2ebb71e9
--- /dev/null
+++ b/tests/integration/system/nsf/interfaces/testbed.h
@@ -0,0 +1,22 @@
+#ifndef PINS_TESTS_INTEGRATION_SYSTEM_NSF_INTERFACES_TESTBED_H_
+#define PINS_TESTS_INTEGRATION_SYSTEM_NSF_INTERFACES_TESTBED_H_
+
+#include <memory>
+#include <variant>
+
+#include "thinkit/generic_testbed.h"
+#include "thinkit/generic_testbed_fixture.h"
+#include "thinkit/mirror_testbed.h"
+#include "thinkit/mirror_testbed_fixture.h"
+
+namespace pins_test {
+
+using Testbed = std::variant<std::unique_ptr<thinkit::GenericTestbed>,
+                             std::unique_ptr<thinkit::MirrorTestbed>>;
+using TestbedInterface =
+    std::variant<std::unique_ptr<thinkit::GenericTestbedInterface>,
+                 std::unique_ptr<thinkit::MirrorTestbedInterface>>;
+
+}  // namespace pins_test
+
+#endif  // PINS_TESTS_INTEGRATION_SYSTEM_NSF_INTERFACES_TESTBED_H_
diff --git a/tests/integration/system/nsf/interfaces/traffic_helper.h b/tests/integration/system/nsf/interfaces/traffic_helper.h
index cac0393b..5ecab806 100644
--- a/tests/integration/system/nsf/interfaces/traffic_helper.h
+++ b/tests/integration/system/nsf/interfaces/traffic_helper.h
@@ -16,7 +16,7 @@
 #define PINS_TESTS_INTEGRATION_SYSTEM_NSF_INTERFACES_TRAFFIC_HELPER_H_
 
 #include "absl/status/status.h"
-#include "thinkit/generic_testbed.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 
 namespace pins_test {
 
@@ -28,15 +28,15 @@ class TrafficHelper {
 
   // Starts traffic with a predefined traffic configuration from a Control
   // Device or Traffic Generator in the testbed.
-  virtual absl::Status StartTraffic(thinkit::GenericTestbed& testbed) = 0;
+  virtual absl::Status StartTraffic(Testbed& testbed) = 0;
 
   // Stops traffic in the testbed.
-  virtual absl::Status StopTraffic(thinkit::GenericTestbed& testbed) = 0;
+  virtual absl::Status StopTraffic(Testbed& testbed) = 0;
 
   // Validates traffic in the testbed.
   // Needs to be called *after* `StopTraffic()` is called.
   virtual absl::Status ValidateTraffic(int error_percentage,
-                                       thinkit::GenericTestbed& testbed) = 0;
+                                       Testbed& testbed) = 0;
 };
 
 }  // namespace pins_test
diff --git a/tests/integration/system/nsf/traffic_helpers/BUILD.bazel b/tests/integration/system/nsf/traffic_helpers/BUILD.bazel
index 97b91642..b8b58476 100644
--- a/tests/integration/system/nsf/traffic_helpers/BUILD.bazel
+++ b/tests/integration/system/nsf/traffic_helpers/BUILD.bazel
@@ -24,8 +24,8 @@ cc_library(
     testonly = True,
     hdrs = ["dvaas_helper.h"],
     deps = [
+        "//tests/integration/system/nsf/interfaces:testbed",
         "//tests/integration/system/nsf/interfaces:traffic_helper",
-        "//thinkit:generic_testbed",
         "@com_google_absl//absl/status",
     ],
 )
@@ -36,10 +36,13 @@ cc_library(
     srcs = ["otg_helper.cc"],
     hdrs = ["otg_helper.h"],
     deps = [
+        "//gutil:overload",
         "//gutil:status",
         "//lib/utils:generic_testbed_utils",
+        "//tests/integration/system/nsf/interfaces:testbed",
         "//tests/integration/system/nsf/interfaces:traffic_helper",
         "//thinkit:generic_testbed",
+        "//thinkit:mirror_testbed",
         "@com_github_grpc_grpc//:grpc++",
         "@com_github_otg_models//:otg_cc_proto",
         "@com_github_otg_models//:otg_grpc_proto",
diff --git a/tests/integration/system/nsf/traffic_helpers/dvaas_helper.h b/tests/integration/system/nsf/traffic_helpers/dvaas_helper.h
index 2a564e90..55fd58a4 100644
--- a/tests/integration/system/nsf/traffic_helpers/dvaas_helper.h
+++ b/tests/integration/system/nsf/traffic_helpers/dvaas_helper.h
@@ -16,21 +16,20 @@
 #define PINS_TESTS_INTEGRATION_SYSTEM_NSF_TRAFFIC_HELPERS_DVAAS_HELPER_H_
 
 #include "absl/status/status.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 #include "tests/integration/system/nsf/interfaces/traffic_helper.h"
-#include "thinkit/generic_testbed.h"
 
 namespace pins_test {
 
 class DvaasHelper : public TrafficHelper {
  public:
-  absl::Status StartTraffic(thinkit::GenericTestbed& testbed) override {
+  absl::Status StartTraffic(Testbed& testbed) override {
     return absl::OkStatus();
   };
-  absl::Status StopTraffic(thinkit::GenericTestbed& testbed) override {
+  absl::Status StopTraffic(Testbed& testbed) override {
     return absl::OkStatus();
   };
-  absl::Status ValidateTraffic(int error_margin,
-                               thinkit::GenericTestbed& testbed) override {
+  absl::Status ValidateTraffic(int error_margin, Testbed& testbed) override {
     return absl::OkStatus();
   };
 };
diff --git a/tests/integration/system/nsf/traffic_helpers/otg_helper.cc b/tests/integration/system/nsf/traffic_helpers/otg_helper.cc
index 0727654f..290b7df3 100644
--- a/tests/integration/system/nsf/traffic_helpers/otg_helper.cc
+++ b/tests/integration/system/nsf/traffic_helpers/otg_helper.cc
@@ -16,7 +16,9 @@
 
 #include <cmath>
 #include <cstdint>
+#include <memory>
 #include <string>
+#include <variant>
 #include <vector>
 
 #include "absl/status/status.h"
@@ -25,175 +27,207 @@
 #include "artifacts/otg.grpc.pb.h"
 #include "artifacts/otg.pb.h"
 #include "grpcpp/client_context.h"
+#include "gutil/overload.h"
 #include "gutil/status.h"
 #include "lib/utils/generic_testbed_utils.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 #include "thinkit/generic_testbed.h"
+#include "thinkit/mirror_testbed.h"
 
 namespace pins_test {
 
-absl::Status OtgHelper::StartTraffic(thinkit::GenericTestbed& testbed) {
-  ASSIGN_OR_RETURN(std::vector<InterfaceLink> up_links,
-                   GetUpLinks(GetAllTrafficGeneratorLinks, testbed));
-  if (up_links.size() <= 2) {
-    return absl::InvalidArgumentError(
-        "Test requires at least 2 SUT ports connected to a Software (Host) or "
-        "Hardware (Ixia) Traffic Generator");
-  }
-
-  // Create config.
-  otg::SetConfigRequest set_config_request;
-  otg::SetConfigResponse set_config_response;
-  grpc::ClientContext set_config_context;
-  auto* config = set_config_request.mutable_config();
-
-  // Randomly pick source and destination hosts.
-  const std::string otg_src_port = up_links.front().peer_interface;
-  const std::string otg_dst_port = up_links.back().peer_interface;
-  const std::string otg_src_mac = up_links.front().peer_mac_address;
-  const std::string otg_dst_mac = up_links.back().peer_mac_address;
-  const std::string otg_src_ip = up_links.front().peer_ipv6_address;
-  const std::string otg_dst_ip = up_links.back().peer_ipv6_address;
-  const std::string otg_src_loc = up_links.front().peer_traffic_location;
-  const std::string otg_dst_loc = up_links.back().peer_traffic_location;
-
-  // Add ports.
-  auto* src_port = config->add_ports();
-  auto* dst_port = config->add_ports();
-  src_port->set_name(otg_src_port);
-  dst_port->set_name(otg_dst_port);
-  src_port->set_location(otg_src_loc);
-  dst_port->set_location(otg_dst_loc);
-
-  // Set layer1.
-  auto* layer1 = config->add_layer1();
-  layer1->set_name("ly");
-  layer1->add_port_names(otg_src_port);
-  layer1->add_port_names(otg_dst_port);
-
-  // Set speed.
-  layer1->set_speed(otg::Layer1::Speed::speed_1_gbps);
-
-  // Set MTU.
-  layer1->set_mtu(9000);
-
-  // Create flow.
-  auto* flow = config->add_flows();
-  flow->set_name("Host Traffic flow");
-
-  // Set Tx / Rx ports.
-  flow->mutable_tx_rx()->set_choice(otg::FlowTxRx::Choice::port);
-  flow->mutable_tx_rx()->mutable_port()->set_tx_name(otg_src_port);
-  flow->mutable_tx_rx()->mutable_port()->set_rx_name(otg_dst_port);
-
-  // Set packet size.
-  flow->mutable_size()->set_choice(otg::FlowSize::Choice::fixed);
-  flow->mutable_size()->set_fixed(256);
-
-  // Set transmission duration.
-  flow->mutable_duration()->set_choice(otg::FlowDuration::Choice::continuous);
-
-  // Set transmission rate.
-  flow->mutable_rate()->set_choice(otg::FlowRate::Choice::percentage);
-  flow->mutable_rate()->set_percentage(10);
-
-  // Set capture metrics.
-  flow->mutable_metrics()->set_enable(true);
-  flow->mutable_metrics()->set_loss(false);
-  flow->mutable_metrics()->set_timestamps(true);
-  flow->mutable_metrics()->mutable_latency()->set_enable(true);
-  flow->mutable_metrics()->mutable_latency()->set_mode(
-      otg::FlowLatencyMetrics::Mode::cut_through);
-
-  // Set ethernet header.
-  auto* eth_packet = flow->add_packet();
-  eth_packet->set_choice(otg::FlowHeader::Choice::ethernet);
-  eth_packet->mutable_ethernet()->mutable_src()->set_choice(
-      otg::PatternFlowEthernetSrc::Choice::value);
-  eth_packet->mutable_ethernet()->mutable_dst()->set_choice(
-      otg::PatternFlowEthernetDst::Choice::value);
-  eth_packet->mutable_ethernet()->mutable_src()->set_value(otg_src_mac);
-  eth_packet->mutable_ethernet()->mutable_dst()->set_value(otg_dst_mac);
-
-  // Set IP header.
-  auto* ip_packet = flow->add_packet();
-  ip_packet->set_choice(otg::FlowHeader::Choice::ipv6);
-  ip_packet->mutable_ipv6()->mutable_src()->set_choice(
-      otg::PatternFlowIpv6Src::Choice::value);
-  ip_packet->mutable_ipv6()->mutable_dst()->set_choice(
-      otg::PatternFlowIpv6Dst::Choice::value);
-  ip_packet->mutable_ipv6()->mutable_src()->set_value(otg_src_ip);
-  ip_packet->mutable_ipv6()->mutable_dst()->set_value(otg_dst_ip);
-
-  // Set the config.
-  otg::Openapi::StubInterface* stub = testbed.GetTrafficClient();
-  RETURN_IF_ERROR(gutil::GrpcStatusToAbslStatus(stub->SetConfig(
-      &set_config_context, set_config_request, &set_config_response)));
-
-  // Start the traffic.
-  otg::SetControlStateRequest request;
-  otg::SetControlStateResponse response;
-  grpc::ClientContext context;
-  request.mutable_control_state()->set_choice(
-      otg::ControlState::Choice::traffic);
-  request.mutable_control_state()->mutable_traffic()->set_choice(
-      otg::StateTraffic::Choice::flow_transmit);
-  request.mutable_control_state()
-      ->mutable_traffic()
-      ->mutable_flow_transmit()
-      ->set_state(otg::StateTrafficFlowTransmit::State::start);
-  return gutil::GrpcStatusToAbslStatus(
-      stub->SetControlState(&context, request, &response));
+absl::Status OtgHelper::StartTraffic(Testbed& testbed) {
+  return std::visit(
+      gutil::Overload{
+          [&](std::unique_ptr<thinkit::GenericTestbed>& testbed)
+              -> absl::Status {
+            ASSIGN_OR_RETURN(std::vector<InterfaceLink> up_links,
+                             GetUpLinks(GetAllTrafficGeneratorLinks, *testbed));
+            if (up_links.size() <= 2) {
+              return absl::InvalidArgumentError(
+                  "Test requires at least 2 SUT ports connected to a Software "
+                  "(Host) or Hardware (Ixia) Traffic Generator");
+            }
+
+            // Create config.
+            otg::SetConfigRequest set_config_request;
+            otg::SetConfigResponse set_config_response;
+            grpc::ClientContext set_config_context;
+            auto* config = set_config_request.mutable_config();
+
+            // Randomly pick source and destination hosts.
+            const std::string otg_src_port = up_links.front().peer_interface;
+            const std::string otg_dst_port = up_links.back().peer_interface;
+            const std::string otg_src_mac = up_links.front().peer_mac_address;
+            const std::string otg_dst_mac = up_links.back().peer_mac_address;
+            const std::string otg_src_ip = up_links.front().peer_ipv6_address;
+            const std::string otg_dst_ip = up_links.back().peer_ipv6_address;
+            const std::string otg_src_loc =
+                up_links.front().peer_traffic_location;
+            const std::string otg_dst_loc =
+                up_links.back().peer_traffic_location;
+
+            // Add ports.
+            auto* src_port = config->add_ports();
+            auto* dst_port = config->add_ports();
+            src_port->set_name(otg_src_port);
+            dst_port->set_name(otg_dst_port);
+            src_port->set_location(otg_src_loc);
+            dst_port->set_location(otg_dst_loc);
+
+            // TODO (b/299256787): Move each of the below configurations into a
+            // helper function. Set layer1.
+            auto* layer1 = config->add_layer1();
+            layer1->set_name("ly");
+            layer1->add_port_names(otg_src_port);
+            layer1->add_port_names(otg_dst_port);
+
+            // Set speed.
+            layer1->set_speed(otg::Layer1::Speed::speed_1_gbps);
+
+            // Set MTU.
+            layer1->set_mtu(9000);
+
+            // Create flow.
+            auto* flow = config->add_flows();
+            flow->set_name("Host Traffic flow");
+
+            // Set Tx / Rx ports.
+            flow->mutable_tx_rx()->set_choice(otg::FlowTxRx::Choice::port);
+            flow->mutable_tx_rx()->mutable_port()->set_tx_name(otg_src_port);
+            flow->mutable_tx_rx()->mutable_port()->set_rx_name(otg_dst_port);
+
+            // Set packet size.
+            flow->mutable_size()->set_choice(otg::FlowSize::Choice::fixed);
+            flow->mutable_size()->set_fixed(256);
+
+            // Set transmission duration.
+            flow->mutable_duration()->set_choice(
+                otg::FlowDuration::Choice::continuous);
+
+            // Set transmission rate.
+            flow->mutable_rate()->set_choice(otg::FlowRate::Choice::percentage);
+            flow->mutable_rate()->set_percentage(10);
+
+            // Set capture metrics.
+            flow->mutable_metrics()->set_enable(true);
+            flow->mutable_metrics()->set_loss(false);
+            flow->mutable_metrics()->set_timestamps(true);
+            flow->mutable_metrics()->mutable_latency()->set_enable(true);
+            flow->mutable_metrics()->mutable_latency()->set_mode(
+                otg::FlowLatencyMetrics::Mode::cut_through);
+
+            // Set ethernet header.
+            auto* eth_packet = flow->add_packet();
+            eth_packet->set_choice(otg::FlowHeader::Choice::ethernet);
+            eth_packet->mutable_ethernet()->mutable_src()->set_choice(
+                otg::PatternFlowEthernetSrc::Choice::value);
+            eth_packet->mutable_ethernet()->mutable_dst()->set_choice(
+                otg::PatternFlowEthernetDst::Choice::value);
+            eth_packet->mutable_ethernet()->mutable_src()->set_value(
+                otg_src_mac);
+            eth_packet->mutable_ethernet()->mutable_dst()->set_value(
+                otg_dst_mac);
+
+            // Set IP header.
+            auto* ip_packet = flow->add_packet();
+            ip_packet->set_choice(otg::FlowHeader::Choice::ipv6);
+            ip_packet->mutable_ipv6()->mutable_src()->set_choice(
+                otg::PatternFlowIpv6Src::Choice::value);
+            ip_packet->mutable_ipv6()->mutable_dst()->set_choice(
+                otg::PatternFlowIpv6Dst::Choice::value);
+            ip_packet->mutable_ipv6()->mutable_src()->set_value(otg_src_ip);
+            ip_packet->mutable_ipv6()->mutable_dst()->set_value(otg_dst_ip);
+
+            // Set the config.
+            otg::Openapi::StubInterface* stub = testbed->GetTrafficClient();
+            RETURN_IF_ERROR(gutil::GrpcStatusToAbslStatus(
+                stub->SetConfig(&set_config_context, set_config_request,
+                                &set_config_response)));
+
+            // Start the traffic.
+            otg::SetControlStateRequest request;
+            otg::SetControlStateResponse response;
+            grpc::ClientContext context;
+            request.mutable_control_state()->set_choice(
+                otg::ControlState::Choice::traffic);
+            request.mutable_control_state()->mutable_traffic()->set_choice(
+                otg::StateTraffic::Choice::flow_transmit);
+            request.mutable_control_state()
+                ->mutable_traffic()
+                ->mutable_flow_transmit()
+                ->set_state(otg::StateTrafficFlowTransmit::State::start);
+            return gutil::GrpcStatusToAbslStatus(
+                stub->SetControlState(&context, request, &response));
+          },
+          [&](std::unique_ptr<thinkit::MirrorTestbed>& testbed) {
+            return absl::UnimplementedError("MirrorTestbed not implemented");
+          }},
+      testbed);
 }
 
-absl::Status OtgHelper::StopTraffic(thinkit::GenericTestbed& testbed) {
+absl::Status OtgHelper::StopTraffic(Testbed& testbed) {
   return absl::UnimplementedError("Stopping traffic is not implemented.");
 }
 
-absl::Status OtgHelper::ValidateTraffic(int error_margin,
-                                        thinkit::GenericTestbed& testbed) {
-  otg::Openapi::StubInterface* stub = testbed.GetTrafficClient();
-  otg::GetMetricsRequest metrics_req;
-  otg::GetMetricsResponse metrics_res;
-  grpc::ClientContext metrics_ctx;
-
-  metrics_req.mutable_metrics_request()->set_choice(
-      otg::MetricsRequest::Choice::flow);
-  RETURN_IF_ERROR(gutil::GrpcStatusToAbslStatus(
-      stub->GetMetrics(&metrics_ctx, metrics_req, &metrics_res)));
-
-  // Verify flow metrics is not empty.
-  if (metrics_res.metrics_response().flow_metrics().empty()) {
-    return absl::InternalError(
-        "Cannot validate traffic as no flow metrics received.");
-  }
-
-  // Verify transmission completion and no frames or bytes drops for each flow.
-  std::vector<std::string> errors;
-  for (const auto& flow_metric :
-       metrics_res.metrics_response().flow_metrics()) {
-    uint64_t bytes_drop = flow_metric.bytes_tx() - flow_metric.bytes_rx();
-    uint64_t frames_drop = flow_metric.frames_tx() - flow_metric.frames_rx();
-    uint64_t bytes_drop_percent = (bytes_drop / flow_metric.bytes_tx()) * 100;
-    float_t frames_drop_percent = (frames_drop / flow_metric.frames_tx()) * 100;
-    bool transmission_stopped =
-        flow_metric.transmit() == otg::FlowMetric::Transmit::stopped;
-
-    if (bytes_drop_percent <= error_margin &&
-        frames_drop_percent <= error_margin)
-      continue;
-
-    errors.push_back(absl::StrCat(
-        "Flow name:\t\t", flow_metric.name(), "\nBytes dropped:\t\t",
-        bytes_drop, "\nFrames dropped:\t\t", frames_drop,
-        transmission_stopped
-            ? ""
-            : "\nTransmission not completed within the expected time."));
-  }
-
-  if (errors.empty()) return absl::OkStatus();
-  return absl::InternalError(absl::StrCat(
-      "Following errors were observed while validating traffic:\n\n",
-      absl::StrJoin(errors, "\n\n")));
+absl::Status OtgHelper::ValidateTraffic(int error_margin, Testbed& testbed) {
+  return std::visit(
+      gutil::Overload{
+          [&](std::unique_ptr<thinkit::GenericTestbed>& testbed)
+              -> absl::Status {
+            otg::Openapi::StubInterface* stub = testbed->GetTrafficClient();
+            otg::GetMetricsRequest metrics_req;
+            otg::GetMetricsResponse metrics_res;
+            grpc::ClientContext metrics_ctx;
+
+            metrics_req.mutable_metrics_request()->set_choice(
+                otg::MetricsRequest::Choice::flow);
+            RETURN_IF_ERROR(gutil::GrpcStatusToAbslStatus(
+                stub->GetMetrics(&metrics_ctx, metrics_req, &metrics_res)));
+
+            // Verify flow metrics is not empty.
+            if (metrics_res.metrics_response().flow_metrics().empty()) {
+              return absl::InternalError(
+                  "Cannot validate traffic as no flow metrics received.");
+            }
+
+            // Verify transmission completion and no frames or bytes drops for
+            // each flow.
+            std::vector<std::string> errors;
+            for (const auto& flow_metric :
+                 metrics_res.metrics_response().flow_metrics()) {
+              uint64_t bytes_drop =
+                  flow_metric.bytes_tx() - flow_metric.bytes_rx();
+              uint64_t frames_drop =
+                  flow_metric.frames_tx() - flow_metric.frames_rx();
+              uint64_t bytes_drop_percent =
+                  (bytes_drop / flow_metric.bytes_tx()) * 100;
+              float_t frames_drop_percent =
+                  (frames_drop / flow_metric.frames_tx()) * 100;
+              bool transmission_stopped =
+                  flow_metric.transmit() == otg::FlowMetric::Transmit::stopped;
+
+              if (bytes_drop_percent <= error_margin &&
+                  frames_drop_percent <= error_margin)
+                continue;
+
+              errors.push_back(absl::StrCat(
+                  "Flow name:\t\t", flow_metric.name(), "\nBytes dropped:\t\t",
+                  bytes_drop, "\nFrames dropped:\t\t", frames_drop,
+                  transmission_stopped ? ""
+                                       : "\nTransmission not completed within "
+                                         "the expected time."));
+            }
+
+            if (errors.empty()) return absl::OkStatus();
+            return absl::InternalError(absl::StrCat(
+                "Following errors were observed while validating traffic:\n\n",
+                absl::StrJoin(errors, "\n\n")));
+          },
+          [&](std::unique_ptr<thinkit::MirrorTestbed>& testbed) {
+            return absl::UnimplementedError("MirrorTestbed not implemented");
+          }},
+      testbed);
 }
 
 }  // namespace pins_test
diff --git a/tests/integration/system/nsf/traffic_helpers/otg_helper.h b/tests/integration/system/nsf/traffic_helpers/otg_helper.h
index b1155186..4f29c985 100644
--- a/tests/integration/system/nsf/traffic_helpers/otg_helper.h
+++ b/tests/integration/system/nsf/traffic_helpers/otg_helper.h
@@ -16,17 +16,16 @@
 #define PINS_TESTS_INTEGRATION_SYSTEM_NSF_TRAFFIC_HELPERS_OTG_HELPER_H_
 
 #include "absl/status/status.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 #include "tests/integration/system/nsf/interfaces/traffic_helper.h"
-#include "thinkit/generic_testbed.h"
 
 namespace pins_test {
 
 class OtgHelper : public TrafficHelper {
  public:
-  absl::Status StartTraffic(thinkit::GenericTestbed& testbed) override;
-  absl::Status StopTraffic(thinkit::GenericTestbed& testbed) override;
-  absl::Status ValidateTraffic(int error_margin,
-                               thinkit::GenericTestbed& testbed) override;
+  absl::Status StartTraffic(Testbed& testbed) override;
+  absl::Status StopTraffic(Testbed& testbed) override;
+  absl::Status ValidateTraffic(int error_margin, Testbed& testbed) override;
 };
 
 }  // namespace pins_test
diff --git a/tests/integration/system/nsf/upgrade_test.cc b/tests/integration/system/nsf/upgrade_test.cc
index 81d01ece..5303115b 100644
--- a/tests/integration/system/nsf/upgrade_test.cc
+++ b/tests/integration/system/nsf/upgrade_test.cc
@@ -31,10 +31,10 @@
 #include "p4/v1/p4runtime.pb.h"
 #include "tests/integration/system/nsf/interfaces/component_validator.h"
 #include "tests/integration/system/nsf/interfaces/flow_programmer.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 #include "tests/integration/system/nsf/interfaces/traffic_helper.h"
 #include "tests/integration/system/nsf/milestone.h"
 #include "tests/integration/system/nsf/util.h"
-#include "thinkit/generic_testbed.h"
 #include "thinkit/proto/generic_testbed.pb.h"
 
 ABSL_FLAG(pins_test::NsfMilestone, milestone, pins_test::NsfMilestone::kAll,
@@ -47,7 +47,7 @@ using ::p4::v1::ReadResponse;
 // Since the validation is while the traffic is in progress, error margin needs
 // to be defined.
 constexpr int kErrorPercentage = 1;
-constexpr absl::Duration kTrafficRunDuration = absl::Seconds(600);
+constexpr absl::Duration kTrafficRunDuration = absl::Minutes(15);
 
 void NsfUpgradeTest::SetUp() {
   flow_programmer_ = GetParam().create_flow_programmer();
@@ -64,78 +64,75 @@ void NsfUpgradeTest::TearDown() { TearDownTestbed(testbed_interface_); }
 // duplicate config push)
 absl::Status NsfUpgradeTest::NsfUpgrade(absl::string_view prev_version,
                                         absl::string_view version) {
-  thinkit::GenericTestbed& testbed =
-      *std::get<std::unique_ptr<thinkit::GenericTestbed>>(testbed_);
-
-  RETURN_IF_ERROR(ValidateSutState(prev_version, testbed, *ssh_client_));
+  RETURN_IF_ERROR(ValidateSutState(prev_version, testbed_, *ssh_client_));
   RETURN_IF_ERROR(ValidateComponents(&ComponentValidator::OnInit,
                                      component_validators_, prev_version,
-                                     testbed));
+                                     testbed_));
   RETURN_IF_ERROR(CaptureDbState());
 
   // P4 Snapshot before programming flows and starting the traffic.
-  ReadResponse snapshot1 = TakeP4FlowSnapshot();
+  ReadResponse p4flow_snapshot1 = TakeP4FlowSnapshot();
 
   // Program all the flows.
-  RETURN_IF_ERROR(flow_programmer_->ProgramFlows(testbed));
+  RETURN_IF_ERROR(flow_programmer_->ProgramFlows(testbed_));
   RETURN_IF_ERROR(ValidateComponents(&ComponentValidator::OnFlowProgram,
                                      component_validators_, prev_version,
-                                     testbed));
+                                     testbed_));
 
-  RETURN_IF_ERROR(traffic_helper_->StartTraffic(testbed));
+  RETURN_IF_ERROR(traffic_helper_->StartTraffic(testbed_));
   RETURN_IF_ERROR(ValidateComponents(&ComponentValidator::OnStartTraffic,
                                      component_validators_, prev_version,
-                                     testbed));
+                                     testbed_));
 
   // P4 Snapshot before Upgrade and NSF reboot.
-  ReadResponse snapshot2 = TakeP4FlowSnapshot();
+  ReadResponse p4flow_snapshot2 = TakeP4FlowSnapshot();
 
-  // Perform Upgrade
+  // Copy image to the switch for installation.
   RETURN_IF_ERROR(Upgrade(version));
   RETURN_IF_ERROR(ValidateComponents(&ComponentValidator::OnUpgrade,
-                                     component_validators_, version, testbed));
+                                     component_validators_, version, testbed_));
 
   // Perform NSF Reboot.
-  RETURN_IF_ERROR(NsfReboot(testbed, *ssh_client_));
-  RETURN_IF_ERROR(WaitForReboot(testbed, *ssh_client_));
+  RETURN_IF_ERROR(NsfReboot(testbed_, *ssh_client_));
+  RETURN_IF_ERROR(WaitForReboot(testbed_, *ssh_client_));
 
   // Perform validations after reboot is completed.
   RETURN_IF_ERROR(ValidateComponents(&ComponentValidator::OnNsfReboot,
-                                     component_validators_, version, testbed));
-  RETURN_IF_ERROR(ValidateSutState(version, testbed, *ssh_client_));
+                                     component_validators_, version, testbed_));
+  RETURN_IF_ERROR(ValidateSutState(version, testbed_, *ssh_client_));
   RETURN_IF_ERROR(ValidateDbState());
 
   // P4 Snapshot after upgrade and NSF reboot.
-  ReadResponse snapshot3 = TakeP4FlowSnapshot();
+  ReadResponse p4flow_snapshot3 = TakeP4FlowSnapshot();
 
   // Push the new config and validate.
   RETURN_IF_ERROR(PushConfig(version));
-  RETURN_IF_ERROR(ValidateSutState(version, testbed, *ssh_client_));
+  RETURN_IF_ERROR(ValidateSutState(version, testbed_, *ssh_client_));
   RETURN_IF_ERROR(ValidateComponents(&ComponentValidator::OnConfigPush,
-                                     component_validators_, version, testbed));
+                                     component_validators_, version, testbed_));
 
   // Wait for transmission duration.
   LOG(INFO) << "Wait for " << kTrafficRunDuration << " for transmit completion";
   absl::SleepFor(kTrafficRunDuration);
 
   // Stop and validate traffic
-  RETURN_IF_ERROR(traffic_helper_->StopTraffic(testbed));
+  RETURN_IF_ERROR(traffic_helper_->StopTraffic(testbed_));
 
-  RETURN_IF_ERROR(traffic_helper_->ValidateTraffic(kErrorPercentage, testbed));
+  RETURN_IF_ERROR(traffic_helper_->ValidateTraffic(kErrorPercentage, testbed_));
   RETURN_IF_ERROR(ValidateComponents(&ComponentValidator::OnStopTraffic,
-                                     component_validators_, version, testbed));
+                                     component_validators_, version, testbed_));
 
   // Selectively clear flows (eg. not clearing nexthop entries for host
   // testbeds).
-  RETURN_IF_ERROR(flow_programmer_->ClearFlows(testbed));
+  RETURN_IF_ERROR(flow_programmer_->ClearFlows(testbed_));
 
   RETURN_IF_ERROR(ValidateComponents(&ComponentValidator::OnFlowCleanup,
-                                     component_validators_, version, testbed));
+                                     component_validators_, version, testbed_));
 
-  ReadResponse snapshot4 = TakeP4FlowSnapshot();
+  ReadResponse p4flow_snapshot4 = TakeP4FlowSnapshot();
 
-  RETURN_IF_ERROR(CompareP4FlowSnapshots(snapshot1, snapshot4));
-  return CompareP4FlowSnapshots(snapshot2, snapshot3);
+  RETURN_IF_ERROR(CompareP4FlowSnapshots(p4flow_snapshot1, p4flow_snapshot4));
+  return CompareP4FlowSnapshots(p4flow_snapshot2, p4flow_snapshot3);
 }
 
 TEST_P(NsfUpgradeTest, UpgradeAndReboot) {
diff --git a/tests/integration/system/nsf/util.cc b/tests/integration/system/nsf/util.cc
index e1e9a72e..082062ff 100644
--- a/tests/integration/system/nsf/util.cc
+++ b/tests/integration/system/nsf/util.cc
@@ -18,6 +18,7 @@
 #include <memory>
 #include <string>
 #include <variant>
+#include <vector>
 
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -36,8 +37,10 @@
 #include "lib/validator/validator_lib.h"
 #include "p4/config/v1/p4info.pb.h"
 #include "tests/integration/system/nsf/interfaces/component_validator.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 #include "thinkit/generic_testbed.h"
 #include "thinkit/generic_testbed_fixture.h"
+#include "thinkit/mirror_testbed.h"
 #include "thinkit/mirror_testbed_fixture.h"
 #include "thinkit/ssh_client.h"
 #include "thinkit/switch.h"
@@ -196,30 +199,46 @@ absl::StatusOr<Testbed> GetTestbed(TestbedInterface& testbed_interface) {
           },
           [&](std::unique_ptr<thinkit::MirrorTestbedInterface>& testbed)
               -> absl::StatusOr<Testbed> {
-            return absl::UnimplementedError(
-                "MirrorTestbedInterface not implemented");
+            return absl::UnimplementedError("MirrorTestbed not implemented");
           }},
       testbed_interface);
 }
 
+thinkit::Switch& GetSut(Testbed& testbed) {
+  return std::visit(
+      gutil::Overload{[&](std::unique_ptr<thinkit::GenericTestbed>& testbed)
+                          -> thinkit::Switch& { return testbed->Sut(); },
+                      [&](std::unique_ptr<thinkit::MirrorTestbed>& testbed)
+                          -> thinkit::Switch& { return testbed->Sut(); }},
+      testbed);
+}
+
 absl::Status InstallRebootPushConfig(absl::string_view version) {
   return absl::OkStatus();
 }
 
-absl::Status ValidateSutState(absl::string_view version,
-                              thinkit::GenericTestbed& testbed,
+absl::Status ValidateSutState(absl::string_view version, Testbed& testbed,
                               thinkit::SSHClient& ssh_client) {
-  return RunReadyValidations(testbed.Sut(), ssh_client,
-                             GetConnectedInterfacesForSut(testbed),
-                             /*check_interfaces_state=*/true,
-                             /*with_healthz=*/true);
+  return std::visit(
+      gutil::Overload{
+          [&](std::unique_ptr<thinkit::GenericTestbed>& testbed) {
+            return RunReadyValidations(testbed->Sut(), ssh_client,
+                                       GetConnectedInterfacesForSut(*testbed),
+                                       /*check_interfaces_state=*/true,
+                                       /*with_healthz=*/true);
+          },
+          [&](std::unique_ptr<thinkit::MirrorTestbed>& testbed) {
+            // TODO: Implement RunReadyValidations for Mirror
+            // testbed scenario.
+            return absl::UnimplementedError("MirrorTestbed not implemented");
+          }},
+      testbed);
 }
 
 absl::Status ValidateComponents(
-    absl::Status (ComponentValidator::*validate)(absl::string_view,
-                                                 thinkit::GenericTestbed&),
+    absl::Status (ComponentValidator::*validate)(absl::string_view, Testbed&),
     const absl::Span<const std::unique_ptr<ComponentValidator>> validators,
-    absl::string_view version, thinkit::GenericTestbed& testbed) {
+    absl::string_view version, Testbed& testbed) {
   for (const std::unique_ptr<ComponentValidator>& validator : validators) {
     RETURN_IF_ERROR((std::invoke(validate, validator, version, testbed)));
   }
@@ -228,10 +247,9 @@ absl::Status ValidateComponents(
 
 absl::Status Upgrade(absl::string_view version) { return absl::OkStatus(); }
 
-absl::Status NsfReboot(thinkit::GenericTestbed& testbed,
-                       thinkit::SSHClient& ssh_client) {
-  ASSIGN_OR_RETURN(auto sut_gnoi_system_stub,
-                   testbed.Sut().CreateGnoiSystemStub());
+absl::Status NsfReboot(Testbed& testbed, thinkit::SSHClient& ssh_client) {
+  thinkit::Switch& sut = GetSut(testbed);
+  ASSIGN_OR_RETURN(auto sut_gnoi_system_stub, sut.CreateGnoiSystemStub());
   gnoi::system::RebootRequest gnoi_request;
   gnoi_request.set_method(gnoi::system::RebootMethod::NSF);
   gnoi_request.set_message("Performing NSF Reboot");
@@ -242,14 +260,26 @@ absl::Status NsfReboot(thinkit::GenericTestbed& testbed,
       &grpc_context, gnoi_request, &gnoi_response));
 }
 
-absl::Status WaitForReboot(thinkit::GenericTestbed& testbed,
-                           thinkit::SSHClient& ssh_client) {
+absl::Status WaitForReboot(Testbed& testbed, thinkit::SSHClient& ssh_client) {
   // Wait for switch to go down and come back up.
-  RETURN_IF_ERROR(WaitForSwitchState(testbed.Sut(), SwitchState::kDown,
-                                     kTurnDownTimeout, ssh_client));
+  thinkit::Switch& sut = GetSut(testbed);
 
-  return WaitForSwitchState(testbed.Sut(), SwitchState::kReady, kTurnUpTimeout,
-                            ssh_client, GetConnectedInterfacesForSut(testbed));
+  return std::visit(
+      gutil::Overload{
+          [&](std::unique_ptr<thinkit::GenericTestbed>& testbed)
+              -> absl::Status {
+            RETURN_IF_ERROR(WaitForSwitchState(sut, SwitchState::kDown,
+                                               kTurnDownTimeout, ssh_client));
+
+            return WaitForSwitchState(sut, SwitchState::kReady, kTurnUpTimeout,
+                                      ssh_client,
+                                      GetConnectedInterfacesForSut(*testbed));
+          },
+          [&](std::unique_ptr<thinkit::MirrorTestbed>& testbed)
+              -> absl::Status {
+            return absl::UnimplementedError("MirrorTestbed not implemented");
+          }},
+      testbed);
 }
 
 absl::Status PushConfig(absl::string_view version) { return absl::OkStatus(); }
@@ -261,8 +291,15 @@ absl::Status CompareP4FlowSnapshots(const ReadResponse& a,
   return absl::OkStatus();
 }
 
-absl::Status CaptureDbState() { return absl::OkStatus(); }
+absl::Status CaptureDbState() {
+  // TODO: Implement gNMI state path capture for comparison later.
+  return absl::OkStatus();
+}
 
-absl::Status ValidateDbState() { return absl::OkStatus(); }
+absl::Status ValidateDbState() {
+  // TODO: Implement gNMI state path validation for comparison of
+  // the various state paths before and after NSF Upgrade/Reboot.
+  return absl::OkStatus();
+}
 
 }  // namespace pins_test
diff --git a/tests/integration/system/nsf/util.h b/tests/integration/system/nsf/util.h
index 10d4bba3..3c63dd58 100644
--- a/tests/integration/system/nsf/util.h
+++ b/tests/integration/system/nsf/util.h
@@ -16,7 +16,6 @@
 #define PINS_TESTS_INTEGRATION_SYSTEM_NSF_UTIL_H_
 
 #include <memory>
-#include <variant>
 
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -24,22 +23,16 @@
 #include "absl/types/span.h"
 #include "p4/config/v1/p4info.pb.h"
 #include "tests/integration/system/nsf/interfaces/component_validator.h"
-#include "thinkit/generic_testbed.h"
-#include "thinkit/generic_testbed_fixture.h"
-#include "thinkit/mirror_testbed.h"
-#include "thinkit/mirror_testbed_fixture.h"
+#include "tests/integration/system/nsf/interfaces/testbed.h"
 #include "thinkit/ssh_client.h"
+#include "thinkit/switch.h"
 
 namespace pins_test {
 
-using Testbed = std::variant<std::unique_ptr<thinkit::GenericTestbed>,
-                             std::unique_ptr<thinkit::MirrorTestbed>>;
-using TestbedInterface =
-    std::variant<std::unique_ptr<thinkit::GenericTestbedInterface>,
-                 std::unique_ptr<thinkit::MirrorTestbedInterface>>;
-
 absl::StatusOr<Testbed> GetTestbed(TestbedInterface& testbed_interface);
 
+thinkit::Switch& GetSut(Testbed& testbed);
+
 void SetupTestbed(TestbedInterface& testbed_interface);
 
 void TearDownTestbed(TestbedInterface& testbed_interface);
@@ -48,23 +41,21 @@ absl::Status InstallRebootPushConfig(absl::string_view version);
 
 // Validates P4, gNMI, SSH connections and port status along with validating the
 // Stack version of the SUT.
-absl::Status ValidateSutState(absl::string_view version,
-                              thinkit::GenericTestbed& testbed,
+absl::Status ValidateSutState(absl::string_view version, Testbed& testbed,
                               thinkit::SSHClient& ssh_client);
 
 absl::Status ValidateComponents(
-    absl::Status (ComponentValidator::*validate)(absl::string_view,
-                                                 thinkit::GenericTestbed&),
+    absl::Status (ComponentValidator::*validate)(absl::string_view, Testbed&),
     absl::Span<const std::unique_ptr<ComponentValidator>> validators,
-    absl::string_view version, thinkit::GenericTestbed& testbed);
+    absl::string_view version, Testbed& testbed);
 
+// Performs image copy on the inactive side using gNOI.
+// Note: This doesn't involve a reboot.
 absl::Status Upgrade(absl::string_view version);
 
-absl::Status NsfReboot(thinkit::GenericTestbed& testbed,
-                       thinkit::SSHClient& ssh_client);
+absl::Status NsfReboot(Testbed& testbed, thinkit::SSHClient& ssh_client);
 
-absl::Status WaitForReboot(thinkit::GenericTestbed& testbed,
-                           thinkit::SSHClient& ssh_client);
+absl::Status WaitForReboot(Testbed& testbed, thinkit::SSHClient& ssh_client);
 
 absl::Status PushConfig(absl::string_view version);