From 60c24430c6d4685032d9e351a537eded7da1a35c Mon Sep 17 00:00:00 2001
From: David Chavez <david@dcvz.io>
Date: Mon, 27 Nov 2023 19:58:43 +0100
Subject: [PATCH] chore(tests): allow overriding autographics backend (#1005)

---
 burn-wgpu/src/graphics.rs | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/burn-wgpu/src/graphics.rs b/burn-wgpu/src/graphics.rs
index c3b3c01f1..de2071c12 100644
--- a/burn-wgpu/src/graphics.rs
+++ b/burn-wgpu/src/graphics.rs
@@ -78,9 +78,27 @@ impl GraphicsApi for WebGpu {
 
 impl GraphicsApi for AutoGraphicsApi {
     fn backend() -> wgpu::Backend {
+        // Allow overriding AutoGraphicsApi backend with ENV var in std test environments
+        #[cfg(not(no_std))]
+        #[cfg(test)]
+        if let Ok(backend_str) = std::env::var("AUTO_GRAPHICS_BACKEND") {
+            match backend_str.to_lowercase().as_str() {
+                "metal" => return wgpu::Backend::Metal,
+                "vulkan" => return wgpu::Backend::Vulkan,
+                "dx12" => return wgpu::Backend::Dx12,
+                "opengl" => return wgpu::Backend::Gl,
+                "webgpu" => return wgpu::Backend::BrowserWebGpu,
+                _ => {
+                    eprintln!("Invalid graphics backend specified in GRAPHICS_BACKEND environment variable");
+                    std::process::exit(1);
+                }
+            }
+        }
+
+        // In a no_std environment or if the environment variable is not set
         #[cfg(target_os = "macos")]
         return wgpu::Backend::Metal;
         #[cfg(not(target_os = "macos"))]
-        wgpu::Backend::Vulkan
+        return wgpu::Backend::Vulkan;
     }
 }