3 files changed, 20 insertions, 2 deletions
diff --git a/unittest/conftest.py b/unittest/conftest.py
index 35cc714..253e1d8 100644
--- a/unittest/conftest.py
+++ b/unittest/conftest.py
@@ -44,7 +44,6 @@ def C(request=None):
     cnt = 0
     a = iter(declarations)
     for declaration in a:
-        if 'NK_device_model' in declaration: continue
         if declaration.strip().startswith('NK_C_API'):
             declaration = declaration.replace('NK_C_API', '').strip()
             while ';' not in declaration:
diff --git a/unittest/test_C_API.cpp b/unittest/test_C_API.cpp
index acfadd2..1964738 100644
--- a/unittest/test_C_API.cpp
+++ b/unittest/test_C_API.cpp
@@ -84,4 +84,18 @@ TEST_CASE("multiple devices with ID", "[BASIC]") {
   }
 
   free (string);
-}
-\ No newline at end of file
+}
+
+TEST_CASE("Get device model", "[BASIC]") {
+    NK_logout();
+    NK_device_model model = NK_get_device_model();
+    REQUIRE(model == NK_device_model::NK_DISCONNECTED);
+
+    auto success = NK_login_auto() == 1;
+    REQUIRE(success);
+    model = NK_get_device_model();
+    REQUIRE(model != NK_device_model::NK_DISCONNECTED);
+
+    REQUIRE((model == NK_PRO || model == NK_STORAGE));
+    NK_logout();
+}
diff --git a/unittest/test_pro.py b/unittest/test_pro.py
index fb936f8..1c61399 100644
--- a/unittest/test_pro.py
+++ b/unittest/test_pro.py
@@ -921,3 +921,8 @@ def test_TOTP_codes_from_nitrokeyapp(secret, C):
     lib_at = lambda : bb(oath.totp(secret, period=period))
     print (lib_at())
     assert lib_at() == code_device
+
+
+def test_get_device_model(C):
+    assert C.NK_get_device_model() != 0
+    # assert C.NK_get_device_model() != C.NK_DISCONNECTED
+\ No newline at end of file