116 lines
4.5 KiB
Python
116 lines
4.5 KiB
Python
from django.contrib.auth import get_user_model
|
|
from django.core.files.uploadedfile import SimpleUploadedFile
|
|
from django.db.models.signals import post_save
|
|
from django.test import TestCase
|
|
from rest_framework import status
|
|
from rest_framework.test import APIClient
|
|
|
|
from apps.accounts.models import Organization, Role
|
|
from apps.knowledge.models import RoleRagDocument, TrainingFile, trigger_ingestion
|
|
|
|
User = get_user_model()
|
|
|
|
class KnowledgeApiTests(TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
super().setUpClass()
|
|
post_save.disconnect(trigger_ingestion, sender=TrainingFile)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
post_save.connect(trigger_ingestion, sender=TrainingFile)
|
|
super().tearDownClass()
|
|
|
|
def setUp(self):
|
|
self.client: APIClient = APIClient()
|
|
self.owner = User.objects.create_user(
|
|
email_address='owner-k@example.com',
|
|
password='pass1234',
|
|
first_name='Owner',
|
|
last_name='K',
|
|
date_of_birth='1990-01-01',
|
|
is_manager=True,
|
|
)
|
|
self.member = User.objects.create_user(
|
|
email_address='member-k@example.com',
|
|
password='pass1234',
|
|
first_name='Member',
|
|
last_name='K',
|
|
date_of_birth='1992-02-02',
|
|
)
|
|
|
|
self.org = Organization.objects.create(name='Knowledge API Org', owner=self.owner)
|
|
self.org.members.add(self.owner, self.member)
|
|
self.role = Role.objects.create(name='Researcher', organization=self.org)
|
|
|
|
self.training_file = TrainingFile.objects.create(
|
|
role=self.role,
|
|
uploaded_by=self.owner,
|
|
file=SimpleUploadedFile('doc.txt', b'content', content_type='text/plain'),
|
|
file_name='doc.txt',
|
|
file_size=7,
|
|
file_type='text/plain',
|
|
)
|
|
self.rag_doc = RoleRagDocument.objects.create(
|
|
role=self.role,
|
|
training_file=self.training_file,
|
|
content='chunk body',
|
|
content_hash='b' * 64,
|
|
metadata={'k': 'v'},
|
|
chunk_index=0,
|
|
)
|
|
|
|
def test_training_file_list_path(self):
|
|
self.client.force_authenticate(self.member)
|
|
response = self.client.get('/api/training-file/')
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
|
|
def test_training_file_create_path(self):
|
|
self.client.force_authenticate(self.member)
|
|
uploaded = SimpleUploadedFile('new.txt', b'new body', content_type='text/plain')
|
|
response = self.client.post('/api/training-file/', {
|
|
'role': str(self.role.uuid),
|
|
'description': 'new file',
|
|
'file': uploaded,
|
|
})
|
|
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
|
|
|
|
def test_training_file_retrieve_path(self):
|
|
self.client.force_authenticate(self.member)
|
|
response = self.client.get(f'/api/training-file/{self.training_file.uuid}/')
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
|
|
def test_training_file_update_path(self):
|
|
self.client.force_authenticate(self.owner)
|
|
response = self.client.put(
|
|
f'/api/training-file/{self.training_file.uuid}/',
|
|
{
|
|
'description': 'updated desc',
|
|
'file': SimpleUploadedFile('replace.txt', b'updated', content_type='text/plain'),
|
|
},
|
|
)
|
|
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
|
|
|
|
def test_training_file_partial_update_path(self):
|
|
self.client.force_authenticate(self.owner)
|
|
response = self.client.patch(
|
|
f'/api/training-file/{self.training_file.uuid}/',
|
|
{'description': 'patched desc'},
|
|
format='multipart',
|
|
)
|
|
self.assertIn(response.status_code, (status.HTTP_200_OK, status.HTTP_400_BAD_REQUEST))
|
|
|
|
def test_training_file_destroy_path(self):
|
|
self.client.force_authenticate(self.owner)
|
|
response = self.client.delete(f'/api/training-file/{self.training_file.uuid}/')
|
|
self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
|
|
|
|
def test_role_rag_document_list_path(self):
|
|
self.client.force_authenticate(self.member)
|
|
response = self.client.get('/api/role-rag-document/')
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
|
|
def test_role_rag_document_retrieve_path(self):
|
|
self.client.force_authenticate(self.member)
|
|
response = self.client.get(f'/api/role-rag-document/{self.rag_doc.uuid}/')
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|