add shorthand version for collections.create_reference

tbarbugli · tbarbugli · commit 76fae7a80f2f · 2018-12-04T14:35:06.000+01:00
diff --git a/stream/collections.py b/stream/collections.py
@@ -9,11 +9,18 @@ def __init__(self, client, token):
         self.client = client
         self.token = token
 
-    def create_reference(self, collection_name, id):
-        _id = id
-        if isinstance(id, (dict,)) and id.get("id") is not None:
-            _id = id.get("id")
-        return "SO:%s:%s" % (collection_name, _id)
+    def create_reference(self, collection_name=None, id=None, entry=None):
+        if isinstance(entry, (dict,)):
+            _collection = entry["collection"]
+            _id = entry["id"]
+        elif collection_name is not None and id is not None:
+            _collection = collection_name
+            _id = id
+        else:
+            raise ValueError(
+                "must call with collection_name and id or with entry arguments"
+            )
+        return "SO:%s:%s" % (_collection, _id)
 
     def upsert(self, collection_name, data):
         """
diff --git a/stream/tests/test_client.py b/stream/tests/test_client.py
@@ -1329,7 +1329,7 @@ def test_feed_enrichment_collection(self):
         activity_data = {
             "actor": "mike",
             "verb": "buy",
-            "object": self.c.collections.create_reference("items", entry),
+            "object": self.c.collections.create_reference(entry=entry),
         }
         f.add_activity(activity_data)
         response = f.get()

Original file line number	Diff line number	Diff line change
`@@ -1329,7 +1329,7 @@ def test_feed_enrichment_collection(self):`
`1329`	`1329`	`activity_data = {`
`1330`	`1330`	`"actor": "mike",`
`1331`	`1331`	`"verb": "buy",`
`1332`		`- "object": self.c.collections.create_reference("items", entry),`
	`1332`	`+ "object": self.c.collections.create_reference(entry=entry),`
`1333`	`1333`	`}`
`1334`	`1334`	`f.add_activity(activity_data)`
`1335`	`1335`	`response = f.get()`