From 948f91f35aa83b15251a08f062e82df8bc1f55e4 Mon Sep 17 00:00:00 2001
From: syl20bnr <sylvain.benner@tracel.ai>
Date: Mon, 2 Dec 2024 23:46:55 -0500
Subject: [PATCH] Add --version argument to xtask test command

---
 .github/workflows/ci.yml   |  9 ++++-----
 README.md                  | 18 ++++++++++++++----
 xtask/src/commands/test.rs | 16 +++++++++++++---
 xtask/src/main.rs          |  4 +++-
 4 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 69b6ab8..e381cb1 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -4,7 +4,6 @@ env:
   TYPOS_LINK: "https://github.com/crate-ci/typos/releases/download"
   TYPOS_VERSION: "1.23.4"
   CUBECL_ROCM_PATH: "/opt/rocm"
-  ROCM_VERSION_FEATURE: "rocm_624"
 
 on:
   push:
@@ -50,13 +49,13 @@ jobs:
       - name: Lint
         run: cargo xtask check lint
       # --------------------------------------------------------------------------------
-      - name: Unit Tests
+      - name: Unit Tests (default ROCm version)
         shell: bash
-        run: cargo xtask test --features ${{ env.ROCM_VERSION_FEATURE }} unit
+        run: cargo xtask test unit
       # --------------------------------------------------------------------------------
-      - name: Integration Tests
+      - name: Integration Tests (default ROCm version)
         shell: bash
-        run: cargo xtask test  --features ${{ env.ROCM_VERSION_FEATURE }} integration
+        run: cargo xtask test integration
       # --------------------------------------------------------------------------------
       - name: Documentation Tests
         shell: bash
diff --git a/README.md b/README.md
index 25e6f56..5777e76 100644
--- a/README.md
+++ b/README.md
@@ -62,19 +62,24 @@ Here is a table of the libraries covered by each crate:
 |:---------------|:-----------------|
 | cubecl-hip-sys | hiprtc, amdhip64 |
 
-
 ## Running tests
 
 To run tests you need to first meet the expectations for both `Prerequisites` and `Usage`
 sections.
 
-Then execute the following xtask command and provide the feature that corresponds to your
-ROCm installation. For instance for the version `6.2.4`:
+Then execute the following xtask command. If you want to test against a different version of
+ROCm than the default one, use the `-v <version>`, for instance `-v 6.2.2`:
 
 ```sh
-cargo xtask test --no-default-features --features rocm_622
+# test default ROCm bindings
+cargo xtask test
+# test a specific version
+cargo xtask test -v rocm_622
 ```
 
+Important: always make sure that ROCm environment variable (see Usage) points to a version that matches the
+tested version.
+
 ## Generate bindings for a given version of ROCm
 
 1) To generate the bindings you need to first meet the expectations for both `Prerequisites`
@@ -98,6 +103,11 @@ rocm_624 = []
 
 4) Replace the default feature in the `Cargo.toml` file be the latest one, in this case `rocm_624`.
 
+```toml
+[features]
+default = ["rocm_624"]
+```
+
 5) Add the generated bindings module to the file `crates/cubecl-hip-sys/src/bindings/mod.rs`
 conditionally to the new feature you just declared as well as the re-exports:
 
diff --git a/xtask/src/commands/test.rs b/xtask/src/commands/test.rs
index bf3e73e..df9b96c 100644
--- a/xtask/src/commands/test.rs
+++ b/xtask/src/commands/test.rs
@@ -1,6 +1,16 @@
 use tracel_xtask::prelude::*;
 
-pub(crate) fn handle_command(mut args: TestCmdArgs) -> anyhow::Result<()> {
-    args.no_default_features = true;
-    base_commands::test::handle_command(args)
+#[macros::extend_command_args(TestCmdArgs, Target, TestSubCommand)]
+pub struct CubeClHipTestCmdArgs {
+    #[arg(long, short)]
+    pub version: Option<String>,
+}
+
+pub(crate) fn handle_command(mut args: CubeClHipTestCmdArgs) -> anyhow::Result<()> {
+    if let Some(version) = args.version.clone() {
+        let feature_name = format!("rocm_{}", version.replace(".", ""));
+        args.no_default_features = true;
+        args.features = Some(vec![feature_name]);
+    }
+    base_commands::test::handle_command(args.try_into().unwrap())
 }
diff --git a/xtask/src/main.rs b/xtask/src/main.rs
index d998dbc..365770a 100644
--- a/xtask/src/main.rs
+++ b/xtask/src/main.rs
@@ -6,10 +6,12 @@ extern crate log;
 use std::time::Instant;
 use tracel_xtask::prelude::*;
 
-#[macros::base_commands(Build, Bump, Check, Compile, Doc, Fix, Publish, Test, Validate)]
+#[macros::base_commands(Build, Bump, Check, Compile, Doc, Fix, Publish, Validate)]
 enum Command {
     /// Generate bindings.
     Bindgen(commands::bindgen::BindgenCmdArgs),
+    /// Test bindings.
+    Test(commands::test::CubeClHipTestCmdArgs),
 }
 
 fn main() -> anyhow::Result<()> {