diff --git a/src/modbus.h b/src/modbus.h
index e2b70d5..37d6018 100644
--- a/src/modbus.h
+++ b/src/modbus.h
@@ -178,6 +178,13 @@ int modbus_reply(modbus_t *ctx, const uint8_t *req,
 
 #define MODBUS_GET_HIGH_BYTE(data) ((data >> 8) & 0xFF)
 #define MODBUS_GET_LOW_BYTE(data) (data & 0xFF)
+#define MODBUS_GET_INT32_FROM_INT16(tab_int16, index) ((tab_int16[index] << 16) + tab_int16[index + 1])
+#define MODBUS_GET_INT16_FROM_INT8(tab_int8, index) ((tab_int8[index] << 8) + tab_int8[index + 1])
+#define MODBUS_SET_INT16_TO_INT8(tab_int8, index, value) \
+    do { \
+       tab_int8[index] = value >> 8; \
+       tab_int8[index + 1] = value & 0xFF; \
+    } while (0)
 
 void modbus_set_bits_from_byte(uint8_t *dest, int address, const uint8_t value);
 void modbus_set_bits_from_bytes(uint8_t *dest, int address, unsigned int nb_bits,
diff --git a/tests/unit-test-server.c b/tests/unit-test-server.c
index 3b3d6f9..27cc4b4 100644
--- a/tests/unit-test-server.c
+++ b/tests/unit-test-server.c
@@ -24,9 +24,6 @@
 
 #include "unit-test.h"
 
-/* Copied from modbus-private.h */
-#define HEADER_LENGTH_TCP 7
-
 enum {
     TCP,
     RTU
@@ -110,22 +107,23 @@ int main(int argc, char*argv[])
 
     for (;;) {
         rc = modbus_receive(ctx, -1, query);
-        if (rc > 0) {
-            if (((query[header_length + 3] << 8) +
-                 query[header_length + 4])
+        if (rc == -1) {
+            /* Connection closed by the client or error */
+            break;
+        }
+
+        /* Read holding registers */
+        if (query[header_length] == 0x03) {
+            if (MODBUS_GET_INT16_FROM_INT8(query, header_length + 3)
                 == UT_REGISTERS_NB_POINTS_SPECIAL) {
-                /* Change the number of values (offset
-                   TCP = 6) */
-                query[header_length + 3] = 0;
-                query[header_length + 4] = UT_REGISTERS_NB_POINTS;
+                printf("Set an incorrect number of values\n");
+                MODBUS_SET_INT16_TO_INT8(query, header_length + 3,
+                                         UT_REGISTERS_NB_POINTS);
             }
+        }
 
-            rc = modbus_reply(ctx, query, rc, mb_mapping);
-            if (rc == -1) {
-                return -1;
-            }
-        } else {
-            /* Connection closed by the client or error */
+        rc = modbus_reply(ctx, query, rc, mb_mapping);
+        if (rc == -1) {
             break;
         }
     }