Merge pull request #16 from RavenPack/fix-creating-edge-dataset

Fix creating EDGE dataset without product attribute specified

Author: joserc87

ravenpackapi/examples/edge_entity_reference_filtered.py

@@ -9,6 +9,7 @@
 you get those mappings for the entities you are interested into, without having
 to download all the rest.
 """
+
 import csv
 
 from ravenpackapi import RPApi

ravenpackapi/examples/ravenpack_annotations_batch.py

@@ -3,6 +3,7 @@
 Note: this example requires the "futures" package in python 2.7,
 or python 3.2+
 """
+
 import glob
 import os
 from concurrent.futures.thread import ThreadPoolExecutor

ravenpackapi/models/dataset.py

@@ -65,11 +65,9 @@ def __init__(
 
         self.uuid = uuid
 
-        # If the product is not specified, we assume it is RPA
-        # unless the uuid is specified, in which case we'll ask the server
         self.product = product
-        if self.product is None and self.uuid is None:
-            self.product = "RPA"
+        if self.product is None and self.uuid is None and self.api is not None:
+            self.product = self.api.product
 
         self.name = name
         self.description = description

ravenpackapi/tests/unit/models/test_dataset.py

@@ -5,14 +5,14 @@
 
 class TestDataset:
     def test_create_rpa_dataset(self, fake_api):
-        dataset = Dataset(api=fake_api, name="test")
+        dataset = Dataset(product="rpa", api=fake_api, name="test")
         assert dataset.name == "test"
         assert dataset.api == fake_api
-        assert dataset.product == "RPA"
+        assert dataset.product == "rpa"
         dataset.save()
         assert fake_api.datasets["1"] == {
             "uuid": "1",
-            "product": "RPA",
+            "product": "rpa",
             "product_version": "1.0",
             "name": "test",
         }
@@ -30,13 +30,18 @@ def test_create_edge_dataset(self, fake_api):
             "name": "test",
         }
 
-    def test_create_dataset_with_default_product(self, fake_api):
+    @pytest.mark.parametrize(
+        "product",
+        ["rpa", "edge"],
+    )
+    def test_create_dataset_with_default_product(self, product):
+        fake_api = FakeAPI(product=product)
         dataset = Dataset(api=fake_api, name="test")
-        assert dataset.product == "RPA"
+        assert dataset.product == product
         dataset.save()
         assert fake_api.datasets["1"] == {
             "uuid": "1",
-            "product": "RPA",
+            "product": product,
             "product_version": "1.0",
             "name": "test",
         }
@@ -117,8 +122,9 @@ def fake_api(self):
 
 
 class FakeAPI:
-    def __init__(self):
+    def __init__(self, product="rpa"):
         self.datasets = {}
+        self.product = product
 
     def request(self, endpoint, method="get", json=None):
         if method == "post":