100 lines
3.8 KiB
Python
100 lines
3.8 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 apps.accounts.models import Organization, Role
|
|
from apps.knowledge.models import RoleRagDocument, TrainingFile, trigger_ingestion
|
|
|
|
User = get_user_model()
|
|
|
|
class KnowledgeModelTests(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.user = User.objects.create_user(
|
|
email_address='knowledge@example.com',
|
|
password='pass1234',
|
|
first_name='Knowledge',
|
|
last_name='User',
|
|
date_of_birth='1991-01-01',
|
|
)
|
|
self.org = Organization.objects.create(name='Knowledge Org', owner=self.user)
|
|
self.org.members.add(self.user)
|
|
self.role = Role.objects.create(name='Analyst', organization=self.org)
|
|
|
|
def test_training_file_fields_and_defaults(self):
|
|
uploaded = SimpleUploadedFile('training.txt', b'hello world', content_type='text/plain')
|
|
training_file = TrainingFile.objects.create(
|
|
organization=self.org,
|
|
role=self.role,
|
|
uploaded_by=self.user,
|
|
file=uploaded,
|
|
file_name='training.txt',
|
|
file_size=11,
|
|
file_type='text/plain',
|
|
description='sample',
|
|
)
|
|
|
|
self.assertEqual(training_file.role, self.role)
|
|
self.assertEqual(training_file.organization, self.org)
|
|
self.assertEqual(training_file.uploaded_by, self.user)
|
|
self.assertEqual(training_file.file_name, 'training.txt')
|
|
self.assertEqual(training_file.file_size, 11)
|
|
self.assertEqual(training_file.file_type, 'text/plain')
|
|
self.assertEqual(training_file.description, 'sample')
|
|
self.assertEqual(training_file.status, 'ingesting')
|
|
self.assertFalse(training_file.is_processed)
|
|
|
|
self.assertIsNotNone(training_file.id)
|
|
self.assertIsNotNone(training_file.uuid)
|
|
self.assertIsNotNone(training_file.created_at)
|
|
self.assertIsNotNone(training_file.updated_at)
|
|
|
|
self.assertIn('training.txt (Analyst)', str(training_file))
|
|
|
|
def test_role_rag_document_fields_and_defaults(self):
|
|
uploaded = SimpleUploadedFile('base.txt', b'base', content_type='text/plain')
|
|
training_file = TrainingFile.objects.create(
|
|
organization=self.org,
|
|
role=self.role,
|
|
uploaded_by=self.user,
|
|
file=uploaded,
|
|
file_name='base.txt',
|
|
file_size=4,
|
|
file_type='text/plain',
|
|
)
|
|
document = RoleRagDocument.objects.create(
|
|
organization=self.org,
|
|
role=self.role,
|
|
training_file=training_file,
|
|
content='Chunk content',
|
|
content_hash='a' * 64,
|
|
metadata={'source': 'base.txt'},
|
|
chunk_index=2,
|
|
is_active=True,
|
|
)
|
|
|
|
self.assertEqual(document.role, self.role)
|
|
self.assertEqual(document.organization, self.org)
|
|
self.assertEqual(document.training_file, training_file)
|
|
self.assertEqual(document.content, 'Chunk content')
|
|
self.assertEqual(document.content_hash, 'a' * 64)
|
|
self.assertEqual(document.metadata, {'source': 'base.txt'})
|
|
self.assertEqual(document.chunk_index, 2)
|
|
self.assertTrue(document.is_active)
|
|
|
|
self.assertIsNotNone(document.id)
|
|
self.assertIsNotNone(document.uuid)
|
|
self.assertIsNotNone(document.created_at)
|
|
self.assertIsNotNone(document.updated_at)
|
|
|
|
self.assertEqual(str(document), 'Analyst - Chunk 2')
|