From 56bcf1ffdd33ed3381e9b3c87dbdc98633367af9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marco=20Trevisan=20=28Trevi=C3=B1o=29?= <mail@3v1n0.net>
Date: Sun, 24 Jan 2021 15:43:49 +0100
Subject: [PATCH] virtual-device: Add command to change the number of enroll
 stages

As per this don't use the class value anymore at enroll phase, as it may
differ.
---
 libfprint/drivers/virtual-device.c | 10 +++++++-
 tests/virtual-device.py            | 37 +++++++++++++++++++++++++++++-
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/libfprint/drivers/virtual-device.c b/libfprint/drivers/virtual-device.c
index 03453b07..8fcff49c 100644
--- a/libfprint/drivers/virtual-device.c
+++ b/libfprint/drivers/virtual-device.c
@@ -40,6 +40,7 @@ G_DEFINE_TYPE (FpDeviceVirtualDevice, fpi_device_virtual_device, FP_TYPE_DEVICE)
 #define ERROR_CMD_PREFIX "ERROR "
 #define RETRY_CMD_PREFIX "RETRY "
 #define FINGER_CMD_PREFIX "FINGER "
+#define SET_ENROLL_STAGES_PREFIX "SET_ENROLL_STAGES "
 
 #define LIST_CMD "LIST"
 
@@ -196,6 +197,13 @@ recv_instruction_cb (GObject      *source_object,
           if (self->prints_storage)
             g_hash_table_foreach (self->prints_storage, write_key_to_listener, listener);
         }
+      else if (g_str_has_prefix (cmd, SET_ENROLL_STAGES_PREFIX))
+        {
+          guint stages;
+
+          stages = g_ascii_strtoull (cmd + strlen (SET_ENROLL_STAGES_PREFIX), NULL, 10);
+          fpi_device_set_nr_enroll_stages (FP_DEVICE (self), stages);
+        }
       else
         {
           g_ptr_array_add (self->pending_commands, g_steal_pointer (&cmd));
@@ -376,7 +384,7 @@ dev_enroll (FpDevice *dev)
 
       self->enroll_stages_passed++;
       fpi_device_enroll_progress (dev, self->enroll_stages_passed, print, NULL);
-      if (self->enroll_stages_passed == FP_DEVICE_GET_CLASS (self)->nr_enroll_stages)
+      if (self->enroll_stages_passed == fp_device_get_nr_enroll_stages (FP_DEVICE (self)))
         {
           if (self->prints_storage)
             {
diff --git a/tests/virtual-device.py b/tests/virtual-device.py
index f2399bd5..48b54946 100644
--- a/tests/virtual-device.py
+++ b/tests/virtual-device.py
@@ -94,7 +94,7 @@ class VirtualDevice(unittest.TestCase):
 
     def send_command(self, command, *args):
         self.assertIn(command, ['INSERT', 'REMOVE', 'SCAN', 'ERROR', 'RETRY',
-            'FINGER'])
+            'FINGER', 'SET_ENROLL_STAGES'])
 
         with Connection(self.sockaddr) as con:
             params = ' '.join(str(p) for p in args)
@@ -145,6 +145,7 @@ class VirtualDevice(unittest.TestCase):
                 self.send_command('SCAN', nick)
                 stage += 1
 
+        self.assertEqual(self._enroll_stage, stage)
         self.assertEqual(self._enroll_stage, self.dev.get_nr_enroll_stages())
         self.assertEqual(self.dev.get_finger_status(), FPrint.FingerStatusFlags.NONE)
 
@@ -246,6 +247,40 @@ class VirtualDevice(unittest.TestCase):
 
         self.assertEqual(self.dev.get_finger_status(), FPrint.FingerStatusFlags.NONE)
 
+    def test_change_enroll_stages(self):
+        notified_spec = None
+        def on_stage_changed(dev, spec):
+            nonlocal notified_spec
+            notified_spec = spec
+
+        self.dev.connect('notify::nr-enroll-stages', on_stage_changed)
+
+        notified_spec = None
+        self.send_command('SET_ENROLL_STAGES', 20)
+        self.assertEqual(self.dev.get_nr_enroll_stages(), 20)
+        self.assertEqual(notified_spec.name, 'nr-enroll-stages')
+
+        notified_spec = None
+        self.send_command('SET_ENROLL_STAGES', 1)
+        self.assertEqual(self.dev.get_nr_enroll_stages(), 1)
+        self.assertEqual(notified_spec.name, 'nr-enroll-stages')
+
+        GLib.test_expect_message('libfprint-device',
+            GLib.LogLevelFlags.LEVEL_CRITICAL, '*enroll_stages > 0*')
+        notified_spec = None
+        self.send_command('SET_ENROLL_STAGES', 0)
+        self.assertEqual(self.dev.get_nr_enroll_stages(), 1)
+        self.assertIsNone(notified_spec)
+        GLib.test_assert_expected_messages_internal('libfprint-device',
+            __file__, 0, 'test_change_enroll_stages')
+
+    def test_quick_enroll(self):
+        self.send_command('SET_ENROLL_STAGES', 1)
+        self.assertEqual(self.dev.get_nr_enroll_stages(), 1)
+        matching = self.enroll_print('testprint', FPrint.Finger.LEFT_LITTLE)
+        self.assertEqual(matching.get_username(), 'testuser')
+        self.assertEqual(matching.get_finger(), FPrint.Finger.LEFT_LITTLE)
+
 class VirtualDeviceStorage(VirtualDevice):
 
     def cleanup_device_storage(self):