import base64 from io import BytesIO from unittest.mock import patch from django.contrib.auth import get_user_model from PIL import Image from rest_framework.authtoken.models import Token from rest_framework.test import APIClient, APITestCase from accounts.models import TransporterDriverLink, TransporterProfile, UserRole, Vehicle User = get_user_model() def _jpeg_data_uri(): image_buffer = BytesIO() Image.new('RGB', (8, 8), color='blue').save(image_buffer, format='JPEG') encoded = base64.b64encode(image_buffer.getvalue()).decode() return f'data:image/jpeg;base64,{encoded}' class AvatarApiTests(APITestCase): def setUp(self): self.avatar = _jpeg_data_uri() @patch('api.serializers.create_traccar_device', return_value={'id': 501}) @patch('api.serializers.generate_unique_tc_uid', return_value='AVATARUID1') def test_individual_driver_register_accepts_avatar(self, _mock_uid, _mock_traccar): response = self.client.post( '/api/transporter/register/', { 'email': 'avatar-driver@test.com', 'password': 'secret123', 'first_name': 'Ava', 'last_name': 'Tar', 'account_type': 'DRIVER', 'local': True, 'country_to_country': False, 'avatar': self.avatar, }, format='json', ) self.assertEqual(response.status_code, 201, response.content) profile = User.objects.get(username='avatar-driver@test.com').transporter_profile self.assertTrue(profile.avatar_url) self.assertTrue(profile.avatar_url.startswith('http')) def test_individual_profile_update_accepts_avatar(self): driver = User.objects.create_user( 'avatar-profile@test.com', email='avatar-profile@test.com', password='secret', ) UserRole.objects.create(user=driver, role=UserRole.Role.TRANSPORTER, phone='', language='en') TransporterProfile.objects.create( user=driver, account_type=TransporterProfile.AccountType.DRIVER, local=True, country_to_country=False, ) token, _ = Token.objects.get_or_create(user=driver) client = APIClient() client.credentials(HTTP_AUTHORIZATION=f'Bearer {token.key}') response = client.patch( '/api/transporter/profile/', {'avatar': self.avatar}, format='json', ) self.assertEqual(response.status_code, 200, response.content) data = response.json()['data'] self.assertTrue(data['avatar_url']) self.assertTrue(data['avatar_url'].startswith('http')) driver.transporter_profile.refresh_from_db() self.assertEqual(driver.transporter_profile.avatar_url, data['avatar_url']) def test_vehicle_create_and_update_accept_avatar(self): owner = User.objects.create_user( 'avatar-vehicle@test.com', email='avatar-vehicle@test.com', password='secret', ) UserRole.objects.create(user=owner, role=UserRole.Role.TRANSPORTER, phone='', language='en') TransporterProfile.objects.create( user=owner, account_type=TransporterProfile.AccountType.FLEET_OWNER, company_name='Fleet Co', local=True, country_to_country=False, ) token, _ = Token.objects.get_or_create(user=owner) client = APIClient() client.credentials(HTTP_AUTHORIZATION=f'Bearer {token.key}') create_response = client.post( '/api/transporter/vehicles/', { 'vehicle_type': 'Truck', 'registration_number': 'AVT-001', 'load_capacity': '10000.00', 'avatar': self.avatar, }, format='json', ) self.assertEqual(create_response.status_code, 201, create_response.content) created = create_response.json()['data'] self.assertTrue(created['avatar_url'].startswith('http')) vehicle = Vehicle.objects.get(pk=created['id']) self.assertEqual(vehicle.avatar_url, created['avatar_url']) update_response = client.patch( f'/api/transporter/vehicles/{vehicle.id}/', {'avatar': self.avatar}, format='json', ) self.assertEqual(update_response.status_code, 200, update_response.content) updated = update_response.json()['data'] self.assertTrue(updated['avatar_url'].startswith('http')) vehicle.refresh_from_db() self.assertEqual(vehicle.avatar_url, updated['avatar_url']) class FleetDriverAvatarApiTests(APITestCase): def setUp(self): self.avatar = _jpeg_data_uri() self.fleet_owner = User.objects.create_user( 'fleet-avatar-owner@test.com', email='fleet-avatar-owner@test.com', password='secret', ) UserRole.objects.create( user=self.fleet_owner, role=UserRole.Role.TRANSPORTER, phone='03000000001', language='en', ) TransporterProfile.objects.create( user=self.fleet_owner, account_type=TransporterProfile.AccountType.FLEET_OWNER, company_name='Avatar Fleet', country_to_country=True, local=False, documents_verified=True, tc_id='1', ) token, _ = Token.objects.get_or_create(user=self.fleet_owner) self.client = APIClient() self.client.credentials(HTTP_AUTHORIZATION=f'Bearer {token.key}') @patch('api.serializers.create_traccar_device', return_value={'id': 888}) @patch('api.serializers.generate_unique_tc_uid', return_value='FLEETAVAT1') def test_create_fleet_driver_accepts_avatar(self, _mock_uid, _mock_traccar): response = self.client.post( '/api/transporter/drivers/', { 'email': 'fleet-avatar-driver@test.com', 'password': 'secret123', 'first_name': 'Fleet', 'last_name': 'Avatar', 'avatar': self.avatar, }, format='json', ) self.assertEqual(response.status_code, 201, response.content) data = response.json()['data'] self.assertTrue(data['avatar_url']) self.assertTrue(data['avatar_url'].startswith('http')) driver = User.objects.get(username='fleet-avatar-driver@test.com') self.assertEqual(driver.transporter_profile.avatar_url, data['avatar_url']) @patch('api.serializers.create_traccar_device', return_value={'id': 889}) @patch('api.serializers.generate_unique_tc_uid', return_value='FLEETAVAT2') def test_patch_fleet_driver_accepts_avatar(self, _mock_uid, _mock_traccar): create_response = self.client.post( '/api/transporter/drivers/', { 'email': 'fleet-avatar-patch@test.com', 'password': 'secret123', 'first_name': 'Patch', 'last_name': 'Driver', }, format='json', ) self.assertEqual(create_response.status_code, 201, create_response.content) driver_id = create_response.json()['data']['id'] response = self.client.patch( f'/api/transporter/drivers/{driver_id}/', {'avatar': self.avatar}, format='json', ) self.assertEqual(response.status_code, 200, response.content) data = response.json()['data'] self.assertTrue(data['avatar_url'].startswith('http')) link = TransporterDriverLink.objects.get( transporter=self.fleet_owner, driver_id=driver_id, ) self.assertEqual(link.driver.transporter_profile.avatar_url, data['avatar_url'])