from django.test import TestCase from django.contrib.auth import get_user_model from rest_framework.test import APITestCase, APIClient from rest_framework import status from rest_framework_simplejwt.tokens import RefreshToken from .models import Company from apps.permissions.models import Permission, Role, UserCompanyRole User = get_user_model() class CompanySecurityTestCase(APITestCase): """Test company management security and permissions""" def setUp(self): """Set up test data""" # Create super admin self.super_admin = User.objects.create_user( email='superadmin@test.com', username='superadmin', password='password123', user_type='super_admin' ) self.super_admin.is_superuser = True self.super_admin.is_active = True # Activate the user self.super_admin.save() # Create regular user self.regular_user = User.objects.create_user( email='user@test.com', username='regularuser', password='password123', user_type='user' ) self.regular_user.is_active = True # Activate the user self.regular_user.save() # Create company self.company = Company.objects.create( name='Test Company', phone='+1234567890', office_time='9 AM - 6 PM', timezone='UTC' ) # Add regular user to company self.regular_user.companies.add(self.company) self.regular_user.active_company = self.company self.regular_user.save() # Create permissions self.create_company_perm = Permission.objects.create( name='Create Company', codename='create_company' ) self.edit_company_perm = Permission.objects.create( name='Edit Company', codename='edit_company' ) self.delete_company_perm = Permission.objects.create( name='Delete Company', codename='delete_company' ) self.manage_users_perm = Permission.objects.create( name='Manage Users', codename='manage_users' ) self.view_company_perm = Permission.objects.create( name='View Company', codename='view_company' ) # Create role with permissions self.admin_role = Role.objects.create( name='Company Admin', description='Company administrator', company=self.company ) self.admin_role.permissions.add( self.create_company_perm, self.edit_company_perm, self.delete_company_perm, self.manage_users_perm, self.view_company_perm ) # Assign role to regular user UserCompanyRole.objects.create( user=self.regular_user, company=self.company, role=self.admin_role ) # Setup API clients self.super_client = APIClient() self.user_client = APIClient() # Get tokens super_refresh = RefreshToken.for_user(self.super_admin) user_refresh = RefreshToken.for_user(self.regular_user) self.super_client.credentials(HTTP_AUTHORIZATION=f'Bearer {super_refresh.access_token}') self.user_client.credentials(HTTP_AUTHORIZATION=f'Bearer {user_refresh.access_token}') def test_super_admin_can_create_company(self): """Super admin should be able to create company without permission check""" data = { 'name': 'New Company', 'phone': '+9876543210', 'office_time': '8 AM - 5 PM', 'timezone': 'UTC' } response = self.super_client.post('/api/companies/companies/', data) self.assertEqual(response.status_code, status.HTTP_201_CREATED) self.assertTrue(Company.objects.filter(name='New Company').exists()) def test_regular_user_cannot_create_company_without_permission(self): """Regular user without create_company permission should be blocked""" # Remove create_company permission from role self.admin_role.permissions.remove(self.create_company_perm) data = { 'name': 'Unauthorized Company', 'phone': '+1111111111', 'office_time': '10 AM - 7 PM', 'timezone': 'UTC' } response = self.user_client.post('/api/companies/companies/', data) self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) def test_regular_user_can_create_company_with_permission(self): """Regular user with create_company permission should be allowed""" data = { 'name': 'Authorized Company', 'phone': '+2222222222', 'office_time': '11 AM - 8 PM', 'timezone': 'UTC' } response = self.user_client.post('/api/companies/companies/', data) self.assertEqual(response.status_code, status.HTTP_201_CREATED) self.assertTrue(Company.objects.filter(name='Authorized Company').exists()) def test_super_admin_can_edit_company(self): """Super admin should be able to edit company""" data = {'name': 'Updated Company Name'} response = self.super_client.patch(f'/api/companies/companies/{self.company.id}/', data) self.assertEqual(response.status_code, status.HTTP_200_OK) self.company.refresh_from_db() self.assertEqual(self.company.name, 'Updated Company Name') def test_regular_user_cannot_edit_company_without_permission(self): """Regular user without edit_company permission should be blocked""" # Remove edit_company permission from role self.admin_role.permissions.remove(self.edit_company_perm) data = {'name': 'Unauthorized Edit'} response = self.user_client.patch(f'/api/companies/companies/{self.company.id}/', data) self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) def test_regular_user_can_edit_company_with_permission(self): """Regular user with edit_company permission should be allowed""" data = {'name': 'Authorized Edit'} response = self.user_client.patch(f'/api/companies/companies/{self.company.id}/', data) self.assertEqual(response.status_code, status.HTTP_200_OK) self.company.refresh_from_db() self.assertEqual(self.company.name, 'Authorized Edit') def test_super_admin_can_delete_company(self): """Super admin should be able to delete company""" response = self.super_client.delete(f'/api/companies/companies/{self.company.id}/') self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT) # Company should be soft deleted (is_active=False) self.company.refresh_from_db() self.assertFalse(self.company.is_active) def test_regular_user_cannot_delete_company_without_permission(self): """Regular user without delete_company permission should be blocked""" # Remove delete_company permission from role self.admin_role.permissions.remove(self.delete_company_perm) response = self.user_client.delete(f'/api/companies/companies/{self.company.id}/') self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) # Company should still exist and be active self.company.refresh_from_db() self.assertTrue(self.company.is_active) def test_regular_user_can_delete_company_with_permission(self): """Regular user with delete_company permission should be allowed""" response = self.user_client.delete(f'/api/companies/companies/{self.company.id}/') self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT) # Company should be soft deleted self.company.refresh_from_db() self.assertFalse(self.company.is_active) def test_super_admin_can_add_user_to_company(self): """Super admin should be able to add user to company""" new_user = User.objects.create_user( email='newuser@test.com', username='newuser', password='password123' ) data = {'user_id': new_user.id} response = self.super_client.post(f'/api/companies/companies/{self.company.id}/add_user/', data) self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertTrue(new_user.companies.filter(id=self.company.id).exists()) def test_regular_user_cannot_add_user_without_permission(self): """Regular user without manage_users permission should be blocked""" # Remove manage_users permission from role self.admin_role.permissions.remove(self.manage_users_perm) new_user = User.objects.create_user( email='newuser2@test.com', username='newuser2', password='password123' ) data = {'user_id': new_user.id} response = self.user_client.post(f'/api/companies/companies/{self.company.id}/add_user/', data) self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) self.assertFalse(new_user.companies.filter(id=self.company.id).exists()) def test_regular_user_can_add_user_with_permission(self): """Regular user with manage_users permission should be allowed""" new_user = User.objects.create_user( email='newuser3@test.com', username='newuser3', password='password123' ) data = {'user_id': new_user.id} response = self.user_client.post(f'/api/companies/companies/{self.company.id}/add_user/', data) self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertTrue(new_user.companies.filter(id=self.company.id).exists()) def test_company_filtering_by_user_membership(self): """Users should only see companies they belong to""" # Create another company that user doesn't belong to other_company = Company.objects.create( name='Other Company', phone='+5555555555', office_time='12 PM - 9 PM', timezone='UTC' ) # User should only see their company response = self.user_client.get('/api/companies/companies/') self.assertEqual(response.status_code, status.HTTP_200_OK) company_names = [company['name'] for company in response.data] self.assertIn('Test Company', company_names) self.assertNotIn('Other Company', company_names) def test_super_admin_sees_all_companies(self): """Super admin should see all companies""" # Create another company other_company = Company.objects.create( name='Other Company', phone='+5555555555', office_time='12 PM - 9 PM', timezone='UTC' ) response = self.super_client.get('/api/companies/companies/') self.assertEqual(response.status_code, status.HTTP_200_OK) company_names = [company['name'] for company in response.data] self.assertIn('Test Company', company_names) self.assertIn('Other Company', company_names)