1 files changed, 28 insertions, 8 deletions
diff --git a/indoteknik_api/controllers/api_v1/wishlist.py b/indoteknik_api/controllers/api_v1/wishlist.py
index 61fa960f..9860f40b 100644
--- a/indoteknik_api/controllers/api_v1/wishlist.py
+++ b/indoteknik_api/controllers/api_v1/wishlist.py
@@ -13,15 +13,26 @@ class Wishlist(controller.Controller):
         
         user_id = int(kw.get('user_id', 0))
         product_id = kw.get('product_id', 0)
+        variant_id = kw.get('variant_id', 0)
         limit = int(kw.get('limit', 0))
         offset = int(kw.get('offset', 0))
         query = [('user_id', '=', user_id)]
         if product_id:
             query += [('product_id', '=', int(product_id))]
+        elif variant_id :
+            query += [('variant_id', '=', int(variant_id))]
         wishlists = request.env['website.user.wishlist'].search(query, limit=limit, offset=offset, order='create_date desc')
+
+        # product = [request.env['product.template'].api_single_response(wishlist.product_id) for wishlist in wishlists]
+        product = []
+        for wishlist in wishlists:
+            if wishlist.product_id:
+                product.append(request.env['product.template'].api_single_response(wishlist.product_id))
+            else:
+                product.append(request.env['product.product'].api_single_response(wishlist.variant_id))
         data = {
             'product_total': request.env['website.user.wishlist'].search_count(query),
-            'products': [request.env['product.template'].api_single_response(wishlist.product_id) for wishlist in wishlists]
+            'products': product 
         }
         return self.response(data)
     
@@ -32,22 +43,31 @@ class Wishlist(controller.Controller):
         
         user_id = int(kw.get('user_id', 0))
         product_id = kw.get('product_id', 0)
-        if not product_id:
-            return self.response(code=400, description='product_id is required')
+        variant_id = kw.get('variant_id', 0)
+        if product_id == 0 and variant_id == 0:
+            return self.response(code=400, description='product_id or variant_id is required')
         
         query = [
             ('user_id', '=', user_id),
-            ('product_id', '=', int(product_id))
         ]
+        if product_id :
+            query +=  [('product_id', '=', int(product_id))]
+            params = { 
+                'user_id': user_id,
+                'product_id': product_id
+            }
+        else:
+            query +=  [('variant_id', '=', int(variant_id))]
+            params = { 
+                'user_id': user_id,
+                'variant_id': variant_id
+            }
         wishlist = request.env['website.user.wishlist'].search(query, limit=1)
         result = {}
         if wishlist:
             wishlist.unlink()
             result['id'] = wishlist.id
         else:
-            create = request.env['website.user.wishlist'].create({
-                'user_id': user_id,
-                'product_id': product_id
-            })
+            create = request.env['website.user.wishlist'].create(params)
             result['id'] = create.id
         return self.response(result)