From 220986998800412e3141a6a18bdf74414aa2ba06 Mon Sep 17 00:00:00 2001
From: Vivien Nicolas <vnicolas@apple.com>
Date: Tue, 30 Nov 2021 04:25:16 +0100
Subject: [PATCH] [YAML] Add a mechanism to specify the test timeout for a
 given YAML file (#11861)

---
 examples/chip-tool/commands/tests/TestCommand.h       | 11 ++++++++++-
 .../chip-tool/templates/partials/test_cluster.zapt    |  3 +++
 src/app/zap-templates/common/ClusterTestGeneration.js |  1 +
 3 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/examples/chip-tool/commands/tests/TestCommand.h b/examples/chip-tool/commands/tests/TestCommand.h
index 7e4d32923e6e01..e19d19ed389f3d 100644
--- a/examples/chip-tool/commands/tests/TestCommand.h
+++ b/examples/chip-tool/commands/tests/TestCommand.h
@@ -30,6 +30,8 @@
 #include <type_traits>
 #include <zap-generated/tests/CHIPClustersTest.h>
 
+constexpr uint16_t kTimeoutInSeconds = 30;
+
 class TestCommand : public CHIPCommand
 {
 public:
@@ -39,20 +41,26 @@ class TestCommand : public CHIPCommand
     {
         AddArgument("node-id", 0, UINT64_MAX, &mNodeId);
         AddArgument("delayInMs", 0, UINT64_MAX, &mDelayInMs);
+        AddArgument("timeout", 0, UINT16_MAX, &mTimeout);
         AddArgument("endpoint-id", CHIP_ZCL_ENDPOINT_MIN, CHIP_ZCL_ENDPOINT_MAX, &mEndpointId);
         AddArgument("PICS", &mPICSFilePath);
     }
 
     /////////// CHIPCommand Interface /////////
     CHIP_ERROR RunCommand() override;
-    chip::System::Clock::Timeout GetWaitDuration() const override { return chip::System::Clock::Seconds16(30); }
+    chip::System::Clock::Timeout GetWaitDuration() const override
+    {
+        return chip::System::Clock::Seconds16(mTimeout.HasValue() ? mTimeout.Value() : kTimeoutInSeconds);
+    }
 
     virtual void NextTest() = 0;
 
     /////////// GlobalCommands Interface /////////
     CHIP_ERROR Wait(chip::System::Clock::Timeout ms);
     CHIP_ERROR WaitForMs(uint16_t ms) { return Wait(chip::System::Clock::Milliseconds32(ms)); }
+    CHIP_ERROR WaitForCommissionee();
     CHIP_ERROR Log(const char * message);
+    CHIP_ERROR Prompt(const char * message);
 
 protected:
     ChipDevice * mDevice;
@@ -261,5 +269,6 @@ class TestCommand : public CHIPCommand
     chip::Optional<uint64_t> mDelayInMs;
     chip::Optional<char *> mPICSFilePath;
     chip::Optional<chip::EndpointId> mEndpointId;
+    chip::Optional<uint16_t> mTimeout;
     chip::Optional<std::map<std::string, bool>> PICS;
 };
diff --git a/examples/chip-tool/templates/partials/test_cluster.zapt b/examples/chip-tool/templates/partials/test_cluster.zapt
index 7729459b86e4a0..713cc0effd4442 100644
--- a/examples/chip-tool/templates/partials/test_cluster.zapt
+++ b/examples/chip-tool/templates/partials/test_cluster.zapt
@@ -44,6 +44,9 @@ class {{filename}}: public TestCommand
       }
     }
 
+  {{#if timeout}}
+    chip::System::Clock::Timeout GetWaitDuration() const override { return chip::System::Clock::Seconds16(mTimeout.HasValue() : mTimeout.Value() : {{timeout}}); }
+  {{/if}}
 
   private:
     std::atomic_uint16_t mTestIndex;
diff --git a/src/app/zap-templates/common/ClusterTestGeneration.js b/src/app/zap-templates/common/ClusterTestGeneration.js
index 1694a291d31a8a..519f6976a02712 100644
--- a/src/app/zap-templates/common/ClusterTestGeneration.js
+++ b/src/app/zap-templates/common/ClusterTestGeneration.js
@@ -327,6 +327,7 @@ function parse(filename)
   });
 
   yaml.filename   = filename;
+  yaml.timeout    = yaml.config.timeout;
   yaml.totalTests = yaml.tests.length;
 
   return yaml;