switched my_company and my_products from method to class views

companies/tests.py

@@ -311,6 +311,9 @@ class MyCompanyViewTest(APITestCase):
         self.user.set_password(self.password)
         self.user.save()
 
+    def tearDown(self):
+        self.model.objects.all().delete()
+
     def test_auth_user_gets_data(self):
         # create instance
         user_instances = [self.factory(creator=self.user) for i in range(5)]
@@ -346,7 +349,7 @@ class MyCompanyViewTest(APITestCase):
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         # assert only 2 instances in response
         payload = response.json()
-        self.assertEquals(2, len(payload))
+        self.assertEquals(2, len(payload['results']))
 
     def test_anon_user_cannot_access(self):
         # send in request

companies/views.py

@@ -23,6 +23,8 @@ from back_latienda.permissions import IsCreator
 from utils import woocommerce
 
 
+
+
 class CompanyViewSet(viewsets.ModelViewSet):
     queryset = Company.objects.filter(is_validated=True).order_by('-created')
     serializer_class = CompanySerializer
@@ -155,23 +157,16 @@ class CompanyViewSet(viewsets.ModelViewSet):
         return Response(message)
 
 
-@api_view(['GET',])
-@permission_classes([IsAuthenticated,])
-def my_company(request):
-    limit = request.GET.get('limit')
-    offset = request.GET.get('offset')
-    qs = Company.objects.filter(creator=request.user)
-    company_serializer = CompanySerializer(qs, many=True)
-    data = company_serializer.data
-    # RESULTS PAGINATION
-    if limit is not None and offset is not None:
-        limit = int(limit)
-        offset = int(offset)
-        data = data[offset:(limit+offset)]
-    elif limit is not None:
-        limit = int(limit)
-        data = data[:limit]
-    return Response(data=data)
+class MyCompanyViewSet(viewsets.ModelViewSet):
+    model = Company
+    serializer_class = CompanySerializer
+    permission_classes = [IsAuthenticated]
+
+    def get_queryset(self):
+        return self.model.objects.filter(creator=self.request.user)
+
+    def perform_create(self, serializer):
+        serializer.save(creator=self.request.user)
 
 
 @api_view(['GET',])