From 69a5578b11719e93cb2d538ff9d1785838b0d495 Mon Sep 17 00:00:00 2001
From: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com>
Date: Wed, 29 Jul 2026 21:29:10 +0800
Subject: [PATCH] Strengthen integrated GPU memory tests

Signed-off-by: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com>
---
 tests-unit/comfy_test/model_management_uma_test.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/tests-unit/comfy_test/model_management_uma_test.py b/tests-unit/comfy_test/model_management_uma_test.py
index f89055bd3..7482068ab 100644
--- a/tests-unit/comfy_test/model_management_uma_test.py
+++ b/tests-unit/comfy_test/model_management_uma_test.py
@@ -7,11 +7,14 @@ import torch
 
 from comfy.cli_args import args as cli_args
 
+original_cpu_arg = cli_args.cpu
 if not torch.cuda.is_available():
     cli_args.cpu = True
 
 import comfy.model_management as model_management  # noqa: E402
 
+cli_args.cpu = original_cpu_arg
+
 
 MIXED_TOPOLOGY = {
     "card1": {
@@ -171,7 +174,12 @@ def test_multiple_candidates_without_bdf_match_do_not_fall_back(monkeypatch):
 def test_non_amd_drm_candidate_is_ignored(monkeypatch):
     cards = {
         "card1": MIXED_TOPOLOGY["card1"],
-        "card2": {**MIXED_TOPOLOGY["card2"], "vendor": "0x10de"},
+        "card2": {
+            **MIXED_TOPOLOGY["card2"],
+            "vendor": "0x10de",
+            "vram": 1,
+            "gtt": 2,
+        },
     }
     install_fake_drm(monkeypatch, cards)
     monkeypatch.setattr(