From af2dc200c52ba65089a34e0136db39d374b32dd0 Mon Sep 17 00:00:00 2001
From: "Dunajski, Bartosz" <bartosz.dunajski@intel.com>
Date: Fri, 15 Feb 2019 13:22:19 +0100
Subject: [PATCH] Improve creating HwContextController

Change-Id: If81ec18793a5af7fb58d66bb26a3bc476eaf94e0
Signed-off-by: Dunajski, Bartosz <bartosz.dunajski@intel.com>
---
 .../command_stream_receiver_simulated_common_hw.inl       | 8 ++++++--
 .../aub_command_stream_receiver_1_tests.cpp               | 5 +++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/runtime/command_stream/command_stream_receiver_simulated_common_hw.inl b/runtime/command_stream/command_stream_receiver_simulated_common_hw.inl
index c61014bc14..fcf163a9e8 100644
--- a/runtime/command_stream/command_stream_receiver_simulated_common_hw.inl
+++ b/runtime/command_stream/command_stream_receiver_simulated_common_hw.inl
@@ -95,8 +95,12 @@ void CommandStreamReceiverSimulatedCommonHw<GfxFamily>::setupContext(OsContext &
 
     auto &engineType = osContext.getEngineType();
     if (aubManager && !(engineType.type == lowPriorityGpgpuEngine.type && engineType.id == lowPriorityGpgpuEngine.id)) {
-        hardwareContextController = std::make_unique<HardwareContextController>(*aubManager, osContext,
-                                                                                deviceIndex, engineIndex, flags);
+        if (osContext.getNumDevicesSupported() == 1) {
+            hardwareContextController = std::make_unique<HardwareContextController>(*aubManager, osContext,
+                                                                                    deviceIndex, engineIndex, flags);
+        } else {
+            hardwareContextController = std::make_unique<HardwareContextController>(*aubManager, osContext, engineIndex, flags);
+        }
     }
 }
 
diff --git a/unit_tests/command_stream/aub_command_stream_receiver_1_tests.cpp b/unit_tests/command_stream/aub_command_stream_receiver_1_tests.cpp
index 4df939da97..8f4007a880 100644
--- a/unit_tests/command_stream/aub_command_stream_receiver_1_tests.cpp
+++ b/unit_tests/command_stream/aub_command_stream_receiver_1_tests.cpp
@@ -803,7 +803,7 @@ HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverWhenProcess
     memoryManager->freeGraphicsMemory(gfxImageAllocation);
 }
 
-HWTEST_F(AubCommandStreamReceiverTests, givenOsContextWithMultipleDevicesSupportedWhenSetupIsCalledThenAbort) {
+HWTEST_F(AubCommandStreamReceiverTests, givenOsContextWithMultipleDevicesSupportedWhenSetupIsCalledThenCreateMultipleHardwareContexts) {
     MockAubCenter *mockAubCenter = new MockAubCenter(platformDevices[0], false, "", CommandStreamReceiverType::CSR_AUB);
     mockAubCenter->aubManager = std::make_unique<MockAubManager>();
     pDevice->executionEnvironment->aubCenter.reset(mockAubCenter);
@@ -811,8 +811,9 @@ HWTEST_F(AubCommandStreamReceiverTests, givenOsContextWithMultipleDevicesSupport
     uint32_t numSupportedDevices = 3;
     OsContext osContext(nullptr, 1, numSupportedDevices, {EngineType::ENGINE_RCS, 0}, PreemptionMode::Disabled);
     auto aubCsr = std::make_unique<AUBCommandStreamReceiverHw<FamilyType>>(*platformDevices[0], "", true, *pDevice->executionEnvironment);
+    aubCsr->setupContext(osContext);
 
-    EXPECT_THROW(aubCsr->setupContext(osContext), std::exception);
+    EXPECT_EQ(numSupportedDevices, aubCsr->hardwareContextController->hardwareContexts.size());
 }
 
 HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverWhenGraphicsAllocationTypeIsntNonAubWritableThenWriteMemoryIsAllowed) {