41 lines
1.7 KiB
Python
41 lines
1.7 KiB
Python
from django.test import TestCase
|
|
from django.contrib.auth import get_user_model
|
|
from apps.orgs.models import Organization, Role
|
|
from apps.mlstore.models import AgentModel, Agent, AgentRun, AgentEvent, RoleRagDocument
|
|
|
|
User = get_user_model()
|
|
|
|
|
|
class MLStoreModelTests(TestCase):
|
|
def setUp(self):
|
|
self.user = User.objects.create_user(email_address='user@example.com', password='pass')
|
|
self.manager = User.objects.create_user(email_address='manager@example.com', password='pass', is_manager=True)
|
|
self.org = Organization.objects.create(name='Org', owner=self.manager)
|
|
self.role = Role.objects.create(name='Engineer', organization=self.org)
|
|
self.model = AgentModel.objects.create(name='test-model', version='v1', path='model.gguf')
|
|
self.agent = Agent.objects.create(model=self.model, organization=self.org)
|
|
|
|
def test_agent_model_str(self):
|
|
self.assertEqual(str(self.model), 'test-model')
|
|
|
|
def test_agent_str(self):
|
|
self.assertIn(self.model.name, str(self.agent))
|
|
|
|
def test_agent_run_str(self):
|
|
run = AgentRun.objects.create(agent=self.agent, user=self.user)
|
|
self.assertIn(str(run.uuid), str(run))
|
|
self.assertIn(str(self.agent), str(run))
|
|
|
|
def test_agent_event_str(self):
|
|
run = AgentRun.objects.create(agent=self.agent, user=self.user)
|
|
evt = AgentEvent.objects.create(execution=run, event_type='message', content={'msg': 'hi'})
|
|
self.assertIn('message', str(evt))
|
|
|
|
def test_role_rag_document_str(self):
|
|
doc = RoleRagDocument.objects.create(
|
|
role=self.role,
|
|
content='chunk',
|
|
content_hash='hash',
|
|
chunk_index=0,
|
|
)
|
|
self.assertIn(self.role.name, str(doc))
|