From 0e70fcdc6f8d45b1facf396f9af63504289d19f2 Mon Sep 17 00:00:00 2001
From: tychovrahe <brunam@seznam.cz>
Date: Wed, 10 Jan 2024 14:56:41 +0100
Subject: [PATCH] feat(core): add haptic to prodtest

[no changelog]
---
 core/embed/prodtest/README.md                 | 15 +++++++++
 core/embed/prodtest/main.c                    | 29 +++++++++++++++++
 core/embed/trezorhal/haptic.h                 |  6 ++++
 .../stm32u5/haptic/drv2625/drv2625.c          | 32 +++++++++++++++----
 4 files changed, 75 insertions(+), 7 deletions(-)

diff --git a/core/embed/prodtest/README.md b/core/embed/prodtest/README.md
index 90e3c28d4..35afc352b 100644
--- a/core/embed/prodtest/README.md
+++ b/core/embed/prodtest/README.md
@@ -160,6 +160,21 @@ SBU 1 0
 OK
 ```
 
+
+### HAPTIC
+The `HAPTIC` command allows you to test the functionality of the device's haptic driver.
+It takes one input parameter, representing the duration of the vibration in milliseconds.
+The device only vibrates if there is motor connected to the haptic driver, otherwise the effect needs to be
+measured by an oscilloscope.
+
+Example:
+```
+// runs the driver for 3000 ms
+
+HAPTIC 3000
+OK
+```
+
 ### OTP READ
 The `OTP READ` command is utilized to retrieve a string parameter from the device's OTP memory.
 This string typically contains information identifying the model and production batch of the device.
diff --git a/core/embed/prodtest/main.c b/core/embed/prodtest/main.c
index 07e5ce5c6..2f7a16728 100644
--- a/core/embed/prodtest/main.c
+++ b/core/embed/prodtest/main.c
@@ -48,6 +48,10 @@
 #include "optiga_transport.h"
 #endif
 
+#ifdef USE_HAPTIC
+#include "haptic.h"
+#endif
+
 #ifdef USE_HASH_PROCESSOR
 #include "hash_processor.h"
 #endif
@@ -460,6 +464,24 @@ static void test_sbu(const char *args) {
 }
 #endif
 
+#ifdef USE_HAPTIC
+static void test_haptic(const char *args) {
+  int duration_ms = atoi(args);
+
+  if (duration_ms <= 0) {
+    vcp_println("ERROR HAPTIC DURATION");
+    return;
+  }
+
+  if (haptic_test(duration_ms)) {
+    vcp_println("OK");
+
+  } else {
+    vcp_println("ERROR HAPTIC");
+  }
+}
+#endif
+
 static void test_otp_read(void) {
   uint8_t data[32];
   memzero(data, sizeof(data));
@@ -574,6 +596,9 @@ int main(void) {
 #endif
 #ifdef USE_SBU
   sbu_init();
+#endif
+#ifdef USE_HAPTIC
+  haptic_init();
 #endif
   usb_init_all();
 
@@ -642,6 +667,10 @@ int main(void) {
     } else if (startswith(line, "SBU ")) {
       test_sbu(line + 4);
 #endif
+#ifdef USE_HAPTIC
+    } else if (startswith(line, "HAPTIC ")) {
+      test_haptic(line + 7);
+#endif
 #ifdef USE_OPTIGA
     } else if (startswith(line, "OPTIGAID READ")) {
       optigaid_read();
diff --git a/core/embed/trezorhal/haptic.h b/core/embed/trezorhal/haptic.h
index 4bd51f726..2cd4c2911 100644
--- a/core/embed/trezorhal/haptic.h
+++ b/core/embed/trezorhal/haptic.h
@@ -2,6 +2,9 @@
 #ifndef TREZORHAL_HAPTIC_H
 #define TREZORHAL_HAPTIC_H
 
+#include <stdbool.h>
+#include <stdint.h>
+
 typedef enum {
   HAPTIC_BUTTON_PRESS = 0,
   HAPTIC_ALERT = 1,
@@ -14,6 +17,9 @@ void haptic_init(void);
 // Calibrate haptic driver
 void haptic_calibrate(void);
 
+// Test haptic driver, plays a maximum amplitude for the given duration
+bool haptic_test(uint16_t duration_ms);
+
 // Play haptic effect
 void haptic_play(haptic_effect_t effect);
 
diff --git a/core/embed/trezorhal/stm32u5/haptic/drv2625/drv2625.c b/core/embed/trezorhal/stm32u5/haptic/drv2625/drv2625.c
index 00e48a745..4134cc51e 100644
--- a/core/embed/trezorhal/stm32u5/haptic/drv2625/drv2625.c
+++ b/core/embed/trezorhal/stm32u5/haptic/drv2625/drv2625.c
@@ -1,6 +1,8 @@
 #include "drv2625_lib.h"
 #include "haptic.h"
 
+#include <stdbool.h>
+
 #include STM32_HAL_H
 
 #include "i2c.h"
@@ -76,10 +78,12 @@
 #define PRESS_EFFECT_AMPLITUDE 25
 #define PRESS_EFFECT_DURATION 10
 
-static void set_reg(uint8_t addr, uint8_t value) {
+#define PRODTEST_EFFECT_AMPLITUDE 127
+
+static bool set_reg(uint8_t addr, uint8_t value) {
   uint8_t data[] = {addr, value};
-  i2c_transmit(DRV2625_I2C_INSTANCE, DRV2625_I2C_ADDRESS, data, sizeof(data),
-               1);
+  return i2c_transmit(DRV2625_I2C_INSTANCE, DRV2625_I2C_ADDRESS, data,
+                      sizeof(data), 1) == HAL_OK;
 }
 
 void haptic_calibrate(void) {
@@ -136,19 +140,29 @@ void haptic_init(void) {
   TIM16->BDTR |= TIM_BDTR_MOE;
 }
 
-static void haptic_play_RTP(int8_t amplitude, uint16_t duration_ms) {
-  set_reg(DRV2625_REG_MODE,
-          DRV2625_REG_MODE_RTP | DRV2625_REG_MODE_TRGFUNC_ENABLE);
-  set_reg(DRV2625_REG_RTP, (uint8_t)amplitude);
+static bool haptic_play_RTP(int8_t amplitude, uint16_t duration_ms) {
+  if (!set_reg(DRV2625_REG_MODE,
+               DRV2625_REG_MODE_RTP | DRV2625_REG_MODE_TRGFUNC_ENABLE)) {
+    return false;
+  }
+
+  if (!set_reg(DRV2625_REG_RTP, (uint8_t)amplitude)) {
+    return false;
+  }
 
   if (duration_ms > 6500) {
     duration_ms = 6500;
   }
+  if (duration_ms == 0) {
+    return true;
+  }
 
   TIM16->CNT = 1;
   TIM16->CCR1 = 1;
   TIM16->ARR = duration_ms * 10;
   TIM16->CR1 |= TIM_CR1_CEN;
+
+  return true;
 }
 
 static void haptic_play_lib(drv2625_lib_effect_t effect) {
@@ -173,3 +187,7 @@ void haptic_play(haptic_effect_t effect) {
       break;
   }
 }
+
+bool haptic_test(uint16_t duration_ms) {
+  return haptic_play_RTP(PRODTEST_EFFECT_AMPLITUDE, duration_ms);
+}