v0.1.1 -- Fixed issue with overriding Target attributes while adding …

…to db
bamhm182 · Feb 12, 2022 · 9dff143 · 9dff143
1 parent b4d20ff
commit 9dff143
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/synack/plugins/db.py b/src/synack/plugins/db.py
@@ -55,7 +55,7 @@ def add_organizations(self, targets, session):
                 db_o = Organization(slug=slug)
                 session.add(db_o)
 
-    def add_targets(self, targets):
+    def add_targets(self, targets, **kwargs):
         session = self.Session()
         self.add_organizations(targets, session)
         q = session.query(Target)
@@ -79,6 +79,8 @@ def add_targets(self, targets):
             db_t.is_registered = t.get('isRegistered')
             db_t.is_updated = t.get('isUpdated')
             db_t.last_submitted = t.get('lastSubmitted')
+            for k in kwargs.keys():
+                setattr(db_t, k, kwargs[k])
         session.commit()
         session.close()
 

diff --git a/test/test_db.py b/test/test_db.py
@@ -338,7 +338,7 @@ def test_add_targets(self):
             "category": {"id": 10}
         }]
 
-        self.db.add_targets(targets)
+        self.db.add_targets(targets, is_registered=True)
         self.db.Session.return_value.commit.assert_called_with()
         self.db.Session.return_value.close.assert_called_with()