From c42d0fbe49a384ea416fb16c003dc9bec9b8026c Mon Sep 17 00:00:00 2001
From: Bartosz Dunajski <bartosz.dunajski@intel.com>
Date: Thu, 8 Aug 2024 11:47:56 +0000
Subject: [PATCH] refactor: improve xe mock engine testing

Signed-off-by: Bartosz Dunajski <bartosz.dunajski@intel.com>
---
 .../test/unit_test/os_interface/linux/xe/mock_drm_xe.inl  | 2 +-
 .../os_interface/linux/xe/mock_drm_xe_definitions.inl     | 8 +++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/shared/test/unit_test/os_interface/linux/xe/mock_drm_xe.inl b/shared/test/unit_test/os_interface/linux/xe/mock_drm_xe.inl
index 88bfa4b9d2..1e2a6f7466 100644
--- a/shared/test/unit_test/os_interface/linux/xe/mock_drm_xe.inl
+++ b/shared/test/unit_test/os_interface/linux/xe/mock_drm_xe.inl
@@ -50,7 +50,7 @@ struct DrmMockXe : public DrmMockCustom {
     StackVec<drm_xe_sync, 1> syncInputs;
     StackVec<drm_xe_ext_set_property, 1> execQueueProperties;
     drm_xe_exec_queue_create latestExecQueueCreate = {};
-    drm_xe_engine_class_instance latestQueueEngineClassInstance = {};
+    std::vector<drm_xe_engine_class_instance> latestQueueEngineClassInstances;
 
     int waitUserFenceReturn = 0;
     int execQueueBanPropertyReturn = 0;
diff --git a/shared/test/unit_test/os_interface/linux/xe/mock_drm_xe_definitions.inl b/shared/test/unit_test/os_interface/linux/xe/mock_drm_xe_definitions.inl
index 4fd2106e76..02bacbf64d 100644
--- a/shared/test/unit_test/os_interface/linux/xe/mock_drm_xe_definitions.inl
+++ b/shared/test/unit_test/os_interface/linux/xe/mock_drm_xe_definitions.inl
@@ -152,7 +152,13 @@ int DrmMockXe::ioctl(DrmIoctl request, void *arg) {
     case DrmIoctl::gemContextCreateExt: {
         auto queueCreate = static_cast<drm_xe_exec_queue_create *>(arg);
         latestExecQueueCreate = *queueCreate;
-        latestQueueEngineClassInstance = reinterpret_cast<drm_xe_engine_class_instance *>(queueCreate->instances)[0];
+        latestQueueEngineClassInstances.clear();
+
+        auto instances = reinterpret_cast<drm_xe_engine_class_instance *>(queueCreate->instances);
+
+        for (uint16_t i = 0; i < queueCreate->num_placements; i++) {
+            latestQueueEngineClassInstances.push_back(instances[i]);
+        }
 
         auto extension = queueCreate->extensions;
         while (extension) {