From b4eada1501852572c1cb78cc385635dbde5969b8 Mon Sep 17 00:00:00 2001
From: Marek Zbroch <marek@zbroch.com>
Date: Sat, 6 Mar 2021 19:57:00 +0100
Subject: [PATCH] Fix Issue #93

---
 .../0004_create_onboardingdevice.py           | 20 +++++++++++++++++++
 netbox_onboarding/worker.py                   |  6 +++++-
 2 files changed, 25 insertions(+), 1 deletion(-)
 create mode 100644 netbox_onboarding/migrations/0004_create_onboardingdevice.py

diff --git a/netbox_onboarding/migrations/0004_create_onboardingdevice.py b/netbox_onboarding/migrations/0004_create_onboardingdevice.py
new file mode 100644
index 0000000..e13a048
--- /dev/null
+++ b/netbox_onboarding/migrations/0004_create_onboardingdevice.py
@@ -0,0 +1,20 @@
+from django.db import migrations
+
+
+def create_missing_onboardingdevice(apps, schema_editor):
+    Device = apps.get_model("dcim", "Device")
+    OnboardingDevice = apps.get_model("netbox_onboarding", "OnboardingDevice")
+
+    for device in Device.objects.filter(onboardingdevice__isnull=True):
+        OnboardingDevice.objects.create(device=device)
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("netbox_onboarding", "0003_onboardingtask_change_logging_model"),
+    ]
+
+    operations = [
+        migrations.RunPython(create_missing_onboardingdevice),
+    ]
diff --git a/netbox_onboarding/worker.py b/netbox_onboarding/worker.py
index 623aa80..d5198b9 100644
--- a/netbox_onboarding/worker.py
+++ b/netbox_onboarding/worker.py
@@ -36,7 +36,7 @@
 
 @REQUEST_TIME.time()
 @job("default")
-def onboard_device(task_id, credentials):  # pylint: disable=too-many-statements
+def onboard_device(task_id, credentials):  # pylint: disable=too-many-statements, too-many-branches
     """Process a single OnboardingTask instance."""
     username = credentials.username
     password = credentials.password
@@ -108,6 +108,10 @@ def onboard_device(task_id, credentials):  # pylint: disable=too-many-statements
         ot.save()
         onboarding_status = False
 
+    finally:
+        if onboarded_device and not OnboardingDevice.objects.filter(device=onboarded_device):
+            OnboardingDevice.objects.create(device=onboarded_device)
+
     onboardingtask_results_counter.labels(status=ot.status).inc()
 
     return dict(ok=onboarding_status)