Add vts test for TWT

Bug: 326998176
Test: atest VtsHalWifiStaIfaceTargetTest
Change-Id: I8efae781963d1e3bec26a2fcc73e174ccbf16f2f
diff --git a/wifi/aidl/vts/functional/wifi_sta_iface_aidl_test.cpp b/wifi/aidl/vts/functional/wifi_sta_iface_aidl_test.cpp
index e8e19ae..51b7301 100644
--- a/wifi/aidl/vts/functional/wifi_sta_iface_aidl_test.cpp
+++ b/wifi/aidl/vts/functional/wifi_sta_iface_aidl_test.cpp
@@ -40,6 +40,8 @@
 using aidl::android::hardware::wifi::StaRoamingCapabilities;
 using aidl::android::hardware::wifi::StaRoamingConfig;
 using aidl::android::hardware::wifi::StaRoamingState;
+using aidl::android::hardware::wifi::TwtCapabilities;
+using aidl::android::hardware::wifi::TwtRequest;
 using aidl::android::hardware::wifi::WifiBand;
 using aidl::android::hardware::wifi::WifiDebugRxPacketFateReport;
 using aidl::android::hardware::wifi::WifiDebugTxPacketFateReport;
@@ -51,6 +53,7 @@
         stopWifiService(getInstanceName());
         wifi_sta_iface_ = getWifiStaIface(getInstanceName());
         ASSERT_NE(nullptr, wifi_sta_iface_.get());
+        ASSERT_TRUE(wifi_sta_iface_->getInterfaceVersion(&interface_version_).isOk());
     }
 
     void TearDown() override { stopWifiService(getInstanceName()); }
@@ -69,6 +72,7 @@
     }
 
     std::shared_ptr<IWifiStaIface> wifi_sta_iface_;
+    int interface_version_;
 
     // Checks if the mDNS Offload is supported by any NIC.
     bool isMdnsOffloadPresentInNIC() {
@@ -356,6 +360,91 @@
     }
 }
 
+TEST_P(WifiStaIfaceAidlTest, TwtGetCapabilities) {
+    if (interface_version_ < 2) {
+        GTEST_SKIP() << "TwtGetCapabilities is available as of sta_iface V2";
+    }
+
+    TwtCapabilities twt_capabilities = {};
+    auto status = wifi_sta_iface_->twtGetCapabilities(&twt_capabilities);
+    if (checkStatusCode(&status, WifiStatusCode::ERROR_NOT_SUPPORTED)) {
+        GTEST_SKIP() << "twtGetCapabilities() is not supported by the vendor";
+    }
+    EXPECT_TRUE(status.isOk());
+    if (!twt_capabilities.isTwtRequesterSupported) {
+        GTEST_SKIP() << "TWT is not supported";
+    }
+
+    EXPECT_GT(twt_capabilities.minWakeDurationUs, 0);
+    EXPECT_GT(twt_capabilities.maxWakeDurationUs, 0);
+    EXPECT_GT(twt_capabilities.minWakeIntervalUs, 0);
+    EXPECT_GT(twt_capabilities.maxWakeIntervalUs, 0);
+}
+
+TEST_P(WifiStaIfaceAidlTest, TwtSessionSetup) {
+    if (interface_version_ < 2) {
+        GTEST_SKIP() << "TwtSessionSetup is available as of sta_iface V2";
+    }
+
+    TwtCapabilities twt_capabilities = {};
+    auto status = wifi_sta_iface_->twtGetCapabilities(&twt_capabilities);
+    if (checkStatusCode(&status, WifiStatusCode::ERROR_NOT_SUPPORTED)) {
+        GTEST_SKIP() << "twtGetCapabilities() is not supported by the vendor";
+    }
+    EXPECT_TRUE(status.isOk());
+    if (!twt_capabilities.isTwtRequesterSupported) {
+        GTEST_SKIP() << "TWT is not supported";
+    }
+
+    TwtRequest twtRequest;
+    twtRequest.mloLinkId = 0;
+    twtRequest.minWakeDurationUs = twt_capabilities.minWakeDurationUs;
+    twtRequest.maxWakeDurationUs = twt_capabilities.maxWakeDurationUs;
+    twtRequest.minWakeIntervalUs = twt_capabilities.minWakeIntervalUs;
+    twtRequest.maxWakeIntervalUs = twt_capabilities.maxWakeIntervalUs;
+    EXPECT_TRUE(wifi_sta_iface_->twtSessionSetup(1, twtRequest).isOk());
+}
+
+TEST_P(WifiStaIfaceAidlTest, TwtSessionGetStats) {
+    if (interface_version_ < 2) {
+        GTEST_SKIP() << "TwtSessionGetStats is available as of sta_iface V2";
+    }
+
+    TwtCapabilities twt_capabilities = {};
+    auto status = wifi_sta_iface_->twtGetCapabilities(&twt_capabilities);
+    if (checkStatusCode(&status, WifiStatusCode::ERROR_NOT_SUPPORTED)) {
+        GTEST_SKIP() << "twtGetCapabilities() is not supported by the vendor";
+    }
+    EXPECT_TRUE(status.isOk());
+    if (!twt_capabilities.isTwtRequesterSupported) {
+        GTEST_SKIP() << "TWT is not supported";
+    }
+
+    // Expecting a IWifiStaIfaceEventCallback.onTwtFailure() with INVALID_PARAMS
+    // as the error code.
+    EXPECT_TRUE(wifi_sta_iface_->twtSessionGetStats(1, 10).isOk());
+}
+
+TEST_P(WifiStaIfaceAidlTest, TwtSessionTeardown) {
+    if (interface_version_ < 2) {
+        GTEST_SKIP() << "TwtSessionTeardown is available as of sta_iface V3";
+    }
+
+    TwtCapabilities twt_capabilities = {};
+    auto status = wifi_sta_iface_->twtGetCapabilities(&twt_capabilities);
+    if (checkStatusCode(&status, WifiStatusCode::ERROR_NOT_SUPPORTED)) {
+        GTEST_SKIP() << "twtGetCapabilities() is not supported by the vendor";
+    }
+    EXPECT_TRUE(status.isOk());
+    if (!twt_capabilities.isTwtRequesterSupported) {
+        GTEST_SKIP() << "TWT is not supported";
+    }
+
+    // Expecting a IWifiStaIfaceEventCallback.onTwtFailure() with INVALID_PARAMS
+    // as  the error code.
+    EXPECT_TRUE(wifi_sta_iface_->twtSessionTeardown(1, 10).isOk());
+}
+
 GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(WifiStaIfaceAidlTest);
 INSTANTIATE_TEST_SUITE_P(WifiTest, WifiStaIfaceAidlTest,
                          testing::ValuesIn(android::getAidlHalInstanceNames(IWifi::descriptor)),