From fc297ee5f8328d4d36921088c62dc925c1b7f444 Mon Sep 17 00:00:00 2001
From: Cayo Puigdefabregas <cayo@ribaguifi.com>
Date: Mon, 17 Oct 2022 18:07:16 +0200
Subject: [PATCH] add is_server_erase in snapshot action

---
 .../d65745749e34_add_is_server_erase.py       | 34 +++++++++++++++++++
 ereuse_devicehub/resources/action/models.py   |  1 +
 .../resources/action/views/snapshot.py        |  7 ++++
 tests/test_render_2_0.py                      | 33 ++++++++++++++++++
 4 files changed, 75 insertions(+)
 create mode 100644 ereuse_devicehub/migrations/versions/d65745749e34_add_is_server_erase.py

diff --git a/ereuse_devicehub/migrations/versions/d65745749e34_add_is_server_erase.py b/ereuse_devicehub/migrations/versions/d65745749e34_add_is_server_erase.py
new file mode 100644
index 00000000..20d001db
--- /dev/null
+++ b/ereuse_devicehub/migrations/versions/d65745749e34_add_is_server_erase.py
@@ -0,0 +1,34 @@
+"""add is_server_erase
+
+Revision ID: d65745749e34
+Revises: a13ed6ad0e3e
+Create Date: 2022-10-17 13:20:29.875274
+
+"""
+import sqlalchemy as sa
+from alembic import context, op
+
+# revision identifiers, used by Alembic.
+revision = 'd65745749e34'
+down_revision = 'a13ed6ad0e3e'
+branch_labels = None
+depends_on = None
+
+
+def get_inv():
+    INV = context.get_x_argument(as_dictionary=True).get('inventory')
+    if not INV:
+        raise ValueError("Inventory value is not specified")
+    return INV
+
+
+def upgrade():
+    op.add_column(
+        'snapshot',
+        sa.Column('is_server_erase', sa.Boolean(), nullable=True),
+        schema=f'{get_inv()}',
+    )
+
+
+def downgrade():
+    op.drop_column('snapshot', 'is_server_erase', schema=f'{get_inv()}')
diff --git a/ereuse_devicehub/resources/action/models.py b/ereuse_devicehub/resources/action/models.py
index 496eed34..87a9d392 100644
--- a/ereuse_devicehub/resources/action/models.py
+++ b/ereuse_devicehub/resources/action/models.py
@@ -676,6 +676,7 @@ class Snapshot(JoinedWithOneDeviceMixin, ActionWithOneDevice):
     of time it took to complete.
     """
     sid = Column(CIText(), nullable=True)
+    is_server_erase = Column(Boolean(), nullable=True)
 
     def get_last_lifetimes(self):
         """We get the lifetime and serial_number of the first disk"""
diff --git a/ereuse_devicehub/resources/action/views/snapshot.py b/ereuse_devicehub/resources/action/views/snapshot.py
index bca48f84..dd967fee 100644
--- a/ereuse_devicehub/resources/action/views/snapshot.py
+++ b/ereuse_devicehub/resources/action/views/snapshot.py
@@ -115,8 +115,15 @@ class SnapshotMixin:
         if snapshot.device.hid is None:
             snapshot.severity = Severity.Warning
 
+        self.is_server_erase(snapshot)
+
         return snapshot
 
+    def is_server_erase(self, snapshot):
+        if snapshot.device.binding:
+            if snapshot.device.binding.kangaroo:
+                snapshot.is_server_erase = True
+
     def get_old_smbios_version(self, debug):
         capabilities = debug.get('lshw', {}).get('capabilities', {})
         for x in capabilities.values():
diff --git a/tests/test_render_2_0.py b/tests/test_render_2_0.py
index 2dcd7100..dbe3c851 100644
--- a/tests/test_render_2_0.py
+++ b/tests/test_render_2_0.py
@@ -2534,3 +2534,36 @@ def test_filter_hdd_in_kangaroo(user3: UserClientFlask):
     assert status == '200 OK'
     for hdd in Device.query.filter_by(type='HardDrive').all():
         assert hdd.dhid in body
+
+
+@pytest.mark.mvp
+@pytest.mark.usefixtures(conftest.app_context.__name__)
+def test_snapshot_is_server_erase(user3: UserClientFlask):
+    snapshot = create_device(user3, 'real-eee-1001pxd.snapshot.12.json')
+
+    user3.get('/workbench/')
+    data = {
+        'csrf_token': generate_csrf(),
+        'phid': snapshot.device.phid(),
+    }
+    user3.post('/workbench/', data=data)
+
+    uri = '/inventory/upload-snapshot/'
+    file_name = 'real-eee-1001pxd.snapshot.12'
+    snapshot_json = conftest.yaml2json(file_name)
+    snapshot_json['uuid'] = 'c058e8d2-fb92-47cb-a4b7-522b75561136'
+    b_snapshot = bytes(json.dumps(snapshot_json), 'utf-8')
+    file_snap = (BytesIO(b_snapshot), file_name)
+    user3.get(uri)
+
+    data = {
+        'snapshot': file_snap,
+        'csrf_token': generate_csrf(),
+    }
+    user3.post(uri, data=data, content_type="multipart/form-data")
+    snapshot2 = Snapshot.query.filter_by(uuid=snapshot_json['uuid']).one()
+
+    assert not snapshot.is_server_erase
+    assert snapshot2.is_server_erase
+    assert snapshot in snapshot.device.actions
+    assert snapshot2 in snapshot.device.actions