feat: add generation source handling for task creation and updates

backend/src/test/kotlin/com/condado/newsletter/controller/EntityTaskControllerTest.kt

@@ -8,6 +8,7 @@ import com.condado.newsletter.scheduler.EntityScheduler
 import com.condado.newsletter.service.JwtService
 import com.ninjasquad.springmockk.MockkBean
 import jakarta.servlet.http.Cookie
+import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.Test
 import org.springframework.beans.factory.annotation.Autowired
@@ -15,6 +16,7 @@ import org.springframework.boot.test.autoconfigure.web.servlet.AutoConfigureMock
 import org.springframework.boot.test.context.SpringBootTest
 import org.springframework.http.MediaType
 import org.springframework.test.web.servlet.MockMvc
+import org.springframework.test.web.servlet.request.MockMvcRequestBuilders.put
 import org.springframework.test.web.servlet.request.MockMvcRequestBuilders.post
 import org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath
 import org.springframework.test.web.servlet.result.MockMvcResultMatchers.status
@@ -56,7 +58,8 @@ class EntityTaskControllerTest {
               "name": "Morning Blast",
               "prompt": "",
               "scheduleCron": "0 8 * * 1-5",
-              "emailLookback": "last_week"
+                            "emailLookback": "last_week",
+                            "generationSource": "openai"
             }
         """.trimIndent()
 
@@ -70,5 +73,96 @@ class EntityTaskControllerTest {
             .andExpect(jsonPath("$.entityId").value(entity.id.toString()))
             .andExpect(jsonPath("$.name").value("Morning Blast"))
             .andExpect(jsonPath("$.prompt").value(""))
+            .andExpect(jsonPath("$.generationSource").value("openai"))
+
+        val persisted = entityTaskRepository.findAll().first()
+        assertThat(persisted.generationSource.value).isEqualTo("openai")
+    }
+
+    @Test
+    fun should_updateTaskAndPersistGenerationSource_when_validRequestProvided() {
+        val entity = virtualEntityRepository.save(
+            VirtualEntity(
+                name = "Entity B",
+                email = "entity-b@condado.com",
+                jobTitle = "Ops"
+            )
+        )
+
+        val createdPayload = """
+            {
+              "entityId": "${entity.id}",
+              "name": "Task One",
+              "prompt": "Initial prompt",
+              "scheduleCron": "0 8 * * 1-5",
+              "emailLookback": "last_week",
+              "generationSource": "openai"
+            }
+        """.trimIndent()
+
+        val createdResult = mockMvc.perform(
+            post("/api/v1/tasks")
+                .cookie(authCookie())
+                .contentType(MediaType.APPLICATION_JSON)
+                .content(createdPayload)
+        )
+            .andExpect(status().isCreated)
+            .andReturn()
+
+        val taskId = com.jayway.jsonpath.JsonPath.read<String>(createdResult.response.contentAsString, "$.id")
+
+        val updatePayload = """
+            {
+              "entityId": "${entity.id}",
+              "name": "Task One Updated",
+              "prompt": "Updated prompt",
+              "scheduleCron": "0 10 * * 1-5",
+              "emailLookback": "last_day",
+              "generationSource": "llama"
+            }
+        """.trimIndent()
+
+        mockMvc.perform(
+            put("/api/v1/tasks/$taskId")
+                .cookie(authCookie())
+                .contentType(MediaType.APPLICATION_JSON)
+                .content(updatePayload)
+        )
+            .andExpect(status().isOk)
+            .andExpect(jsonPath("$.name").value("Task One Updated"))
+            .andExpect(jsonPath("$.generationSource").value("llama"))
+
+        val persisted = entityTaskRepository.findById(java.util.UUID.fromString(taskId)).orElseThrow()
+        assertThat(persisted.generationSource.value).isEqualTo("llama")
+    }
+
+    @Test
+    fun should_returnBadRequest_when_generationSourceIsInvalid() {
+        val entity = virtualEntityRepository.save(
+            VirtualEntity(
+                name = "Entity C",
+                email = "entity-c@condado.com",
+                jobTitle = "Ops"
+            )
+        )
+
+        val payload = """
+            {
+              "entityId": "${entity.id}",
+              "name": "Morning Blast",
+              "prompt": "Prompt",
+              "scheduleCron": "0 8 * * 1-5",
+              "emailLookback": "last_week",
+              "generationSource": "invalid-provider"
+            }
+        """.trimIndent()
+
+        mockMvc.perform(
+            post("/api/v1/tasks")
+                .cookie(authCookie())
+                .contentType(MediaType.APPLICATION_JSON)
+                .content(payload)
+        )
+            .andExpect(status().isBadRequest)
     }
 }

backend/src/test/kotlin/com/condado/newsletter/service/TaskGeneratedMessageServiceTest.kt

@@ -5,6 +5,8 @@ import com.condado.newsletter.dto.TaskPreviewGenerateRequestDto
 import com.condado.newsletter.dto.TaskPreviewTaskDto
 import com.condado.newsletter.model.EntityTask
 import com.condado.newsletter.model.GeneratedMessageHistory
+import com.condado.newsletter.model.ParsedAiResponse
+import com.condado.newsletter.model.TaskGenerationSource
 import com.condado.newsletter.model.VirtualEntity
 import com.condado.newsletter.repository.EntityTaskRepository
 import com.condado.newsletter.repository.GeneratedMessageHistoryRepository
@@ -21,15 +23,17 @@ class TaskGeneratedMessageServiceTest {
     private val generatedMessageHistoryRepository: GeneratedMessageHistoryRepository = mockk()
     private val entityTaskRepository: EntityTaskRepository = mockk()
     private val llamaPreviewService: LlamaPreviewService = mockk()
+    private val aiService: AiService = mockk()
 
     private val service = TaskGeneratedMessageService(
         generatedMessageHistoryRepository = generatedMessageHistoryRepository,
         entityTaskRepository = entityTaskRepository,
-        llamaPreviewService = llamaPreviewService
+        llamaPreviewService = llamaPreviewService,
+        aiService = aiService
     )
 
     @Test
-    fun should_generateAndPersistMessage_when_generateAndSaveCalled() {
+    fun should_useLlamaProvider_when_taskGenerationSourceIsLlama() {
         val taskId = UUID.randomUUID()
         val entity = VirtualEntity(name = "Entity", email = "e@x.com", jobTitle = "Ops").apply { id = UUID.randomUUID() }
         val task = EntityTask(
@@ -37,7 +41,8 @@ class TaskGeneratedMessageServiceTest {
             name = "Task",
             prompt = "Prompt",
             scheduleCron = "0 9 * * 1",
-            emailLookback = "last_week"
+            emailLookback = "last_week",
+            generationSource = TaskGenerationSource.LLAMA
         ).apply { id = taskId }
         val captured = slot<GeneratedMessageHistory>()
 
@@ -59,9 +64,40 @@ class TaskGeneratedMessageServiceTest {
         assertThat(captured.captured.task.id).isEqualTo(taskId)
 
         verify(exactly = 1) { llamaPreviewService.generate(any()) }
+        verify(exactly = 0) { aiService.generate(any()) }
         verify(exactly = 1) { generatedMessageHistoryRepository.save(any()) }
     }
 
+    @Test
+    fun should_useOpenAiProvider_when_taskGenerationSourceIsOpenai() {
+        val taskId = UUID.randomUUID()
+        val entity = VirtualEntity(name = "Entity", email = "e@x.com", jobTitle = "Ops").apply { id = UUID.randomUUID() }
+        val task = EntityTask(
+            virtualEntity = entity,
+            name = "Task",
+            prompt = "Prompt",
+            scheduleCron = "0 9 * * 1",
+            emailLookback = "last_week",
+            generationSource = TaskGenerationSource.OPENAI
+        ).apply { id = taskId }
+        val captured = slot<GeneratedMessageHistory>()
+
+        every { aiService.generate(any()) } returns ParsedAiResponse(subject = "Open Subject", body = "Open Body")
+        every { entityTaskRepository.findById(taskId) } returns java.util.Optional.of(task)
+        every { generatedMessageHistoryRepository.countByTask_Id(taskId) } returns 0
+        every { generatedMessageHistoryRepository.save(capture(captured)) } answers {
+            captured.captured.apply {
+                id = UUID.fromString("00000000-0000-0000-0000-000000000001")
+            }
+        }
+
+        val response = service.generateAndSave(taskId, sampleRequest())
+
+        assertThat(response.content).isEqualTo("SUBJECT: Open Subject\nBODY:\nOpen Body")
+        verify(exactly = 1) { aiService.generate(any()) }
+        verify(exactly = 0) { llamaPreviewService.generate(any()) }
+    }
+
     private fun sampleRequest() = TaskPreviewGenerateRequestDto(
         entity = TaskPreviewEntityDto(
             id = UUID.randomUUID().toString(),