4 files changed, 4 insertions, 8 deletions

diff --git a/tests/basic/test.cpp b/tests/basic/test.cpp
index 26a09585c..1bd5c2762 100644
--- a/tests/basic/test.cpp
+++ b/tests/basic/test.cpp
@@ -37,7 +37,6 @@ class KeyPress : public TestFixture {};
 
 TEST_F(KeyPress, SendKeyboardIsNotCalledWhenNoKeyIsPressed) {
     TestDriver driver;
-    EXPECT_CALL(driver, keyboard_leds_mock()).WillRepeatedly(Return(0));
     EXPECT_CALL(driver, send_keyboard_mock(_)).Times(0);
     keyboard_task();
 }
@@ -45,7 +44,6 @@ TEST_F(KeyPress, SendKeyboardIsNotCalledWhenNoKeyIsPressed) {
 TEST_F(KeyPress, CorrectKeyIsReportedWhenPressed) {
     TestDriver driver;
     press_key(0, 0);
-    EXPECT_CALL(driver, keyboard_leds_mock()).WillRepeatedly(Return(0));
     EXPECT_CALL(driver, send_keyboard_mock(KeyboardReport(KC_A)));
     keyboard_task();
 }
@@ -54,7 +52,6 @@ TEST_F(KeyPress, CorrectKeysAreReportedWhenTwoKeysArePressed) {
     TestDriver driver;
     press_key(1, 0);
     press_key(0, 1);
-    EXPECT_CALL(driver, keyboard_leds_mock()).WillRepeatedly(Return(0));
     //Note that QMK only processes one key at a time
     EXPECT_CALL(driver, send_keyboard_mock(KeyboardReport(KC_B)));
     keyboard_task();
diff --git a/tests/test_common/test_driver.cpp b/tests/test_common/test_driver.cpp
index 9e618aa97..feb80563a 100644
--- a/tests/test_common/test_driver.cpp
+++ b/tests/test_common/test_driver.cpp
@@ -27,7 +27,6 @@ TestDriver::TestDriver()
         &TestDriver::send_consumer
     }
 {
-
     host_set_driver(&m_driver);
     m_this = this;
 }
@@ -37,7 +36,7 @@ TestDriver::~TestDriver() {
 }
 
 uint8_t TestDriver::keyboard_leds(void) {
-    return m_this->keyboard_leds_mock();
+    return m_this->m_leds;
 }
 
 void TestDriver::send_keyboard(report_keyboard_t* report) {
diff --git a/tests/test_common/test_driver.h b/tests/test_common/test_driver.h
index b1b95fbcc..0123fd539 100644
--- a/tests/test_common/test_driver.h
+++ b/tests/test_common/test_driver.h
@@ -27,7 +27,8 @@ class TestDriver {
 public:
     TestDriver();
     ~TestDriver();
-    MOCK_METHOD0(keyboard_leds_mock, uint8_t ());
+    void set_leds(uint8_t leds) { m_leds = leds; }
+    
     MOCK_METHOD1(send_keyboard_mock, void (report_keyboard_t&));
     MOCK_METHOD1(send_mouse_mock, void (report_mouse_t&));
     MOCK_METHOD1(send_system_mock, void (uint16_t));
@@ -39,6 +40,7 @@ private:
     static void send_system(uint16_t data);
     static void send_consumer(uint16_t data);
     host_driver_t m_driver;
+    uint8_t m_leds = 0;
     static TestDriver* m_this;
 };
 
diff --git a/tests/test_common/test_fixture.cpp b/tests/test_common/test_fixture.cpp
index aff518d21..eef9b854b 100644
--- a/tests/test_common/test_fixture.cpp
+++ b/tests/test_common/test_fixture.cpp
@@ -27,12 +27,10 @@ TestFixture::~TestFixture() {
     // Run for a while to make sure all keys are completely released
     // Should probably wait until tapping term etc, has timed out
     EXPECT_CALL(driver, send_keyboard_mock(_)).Times(AnyNumber());
-    EXPECT_CALL(driver, keyboard_leds_mock()).WillRepeatedly(Return(0));
     for (int i=0; i<100; i++) {
         keyboard_task();
     }
     testing::Mock::VerifyAndClearExpectations(&driver); 
     // Verify that the matrix really is cleared
     EXPECT_CALL(driver, send_keyboard_mock(KeyboardReport())).Times(Between(0, 1));
-    EXPECT_CALL(driver, keyboard_leds_mock()).WillRepeatedly(Return(0));
 }
\ No newline at end of file