@@ -33,6 +33,7 @@ import org.junit.jupiter.api.DisplayName
3333import org.junit.jupiter.api.Test
3434import org.junit.jupiter.api.Timeout
3535import org.msgpack.core.MessagePack
36+ import org.msgpack.core.buffer.ArrayBufferInput
3637import java.io.ByteArrayOutputStream
3738import java.util.concurrent.ArrayBlockingQueue
3839import java.util.concurrent.TimeUnit
@@ -60,6 +61,111 @@ class ServerTest {
6061 writeByteArray(payload)
6162 }
6263
64+ private fun errorResponseFrame (id : Int ): ByteArray {
65+ val out = ByteArrayOutputStream ()
66+ MessagePack .newDefaultPacker(out ).use { packer ->
67+ packer.packArrayHeader(3 )
68+ packer.packInt(id)
69+ packer.packNil()
70+ packer.packMapHeader(3 )
71+ packer.packString(" type" )
72+ packer.packString(" ErrorResponse" )
73+ packer.packString(" error" )
74+ packer.packString(" API_SERVER_ERROR" )
75+ packer.packString(" detail" )
76+ packer.packMapHeader(1 )
77+ packer.packString(" status_code" )
78+ packer.packInt(500 )
79+ }
80+ return out .toByteArray()
81+ }
82+
83+ private fun startupFrame (
84+ mapIndex : Int? ,
85+ tiContext : Map <String , Any ?>,
86+ ): ByteArray {
87+ val body =
88+ mapOf (
89+ " type" to " StartupDetails" ,
90+ " ti" to
91+ mapOf (
92+ " id" to " 4d828a62-a417-4936-a7a6-2b3fabacecab" ,
93+ " task_id" to " extract" ,
94+ " dag_id" to " dag1" ,
95+ " run_id" to " run1" ,
96+ " try_number" to 2 ,
97+ " map_index" to mapIndex,
98+ " dag_version_id" to " 4d828a62-a417-4936-a7a6-2b3fabacecab" ,
99+ ),
100+ " ti_context" to mapOf (" dag_run" to mapOf (" dag_id" to " dag1" , " run_id" to " run1" ), " max_tries" to 1 ) + tiContext,
101+ " dag_rel_path" to " /dev/null" ,
102+ " bundle_info" to mapOf (" name" to " any-name" , " version" to " any-version" ),
103+ " start_date" to " 2024-12-01T01:00:00Z" ,
104+ " sentry_integration" to " " ,
105+ )
106+ return Frame .encodeRequest(0 , body).fold(ByteArray (0 )) { acc, buffer -> acc + buffer.toByteArray() }
107+ }
108+
109+ /* *
110+ * Runs [taskClass] as task "extract" of Dag "dag1" through [Server.dispatchTask]. The fake
111+ * supervisor first sends StartupDetails, with [mapIndex] as the map_index of the task instance and
112+ * [tiContext] added to ti_context. Until the task reports its final state, [answer] builds the
113+ * reply to each request from the request ID. Returns the bodies of those requests in arrival
114+ * order, and the body of the final state.
115+ */
116+ private fun serveTask (
117+ taskClass : Class <out Task >,
118+ mapIndex : Int? ,
119+ tiContext : Map <String , Any ?>,
120+ answer : (Int ) -> ByteArray = ::ackFrame,
121+ ): Pair <List <Map <* , * >>, Map<*, *>> {
122+ val toServer = ByteChannel (autoFlush = true )
123+ val fromServer = ByteChannel (autoFlush = true )
124+ val comm = CoordinatorComm (toServer, fromServer)
125+ val server = Server (InetSocketAddress (" localhost" , 0 ), InetSocketAddress (" localhost" , 0 ))
126+ val bundle = Bundle (listOf (DagDef (" dag1" ).addTask(TaskDef (" extract" , taskClass))))
127+
128+ val requests = mutableListOf<Map <* , * >>()
129+ var terminal: Map <* , * >? = null
130+ val supervisor =
131+ Thread {
132+ runBlocking {
133+ toServer.writeFrame(startupFrame(mapIndex, tiContext))
134+ while (terminal == null ) {
135+ val prefix = fromServer.readByteArray(4 )
136+ val payload = fromServer.readByteArray(Frame .parseLengthPrefix(prefix).toInt())
137+ val request = Frame .decodeRaw(ArrayBufferInput (payload))
138+ val body = request.rawBody as Map <* , * >
139+ if (body[" type" ] in setOf (" SucceedTask" , " TaskState" , " RetryTask" )) {
140+ terminal = body
141+ toServer.writeFrame(ackFrame(request.id))
142+ } else {
143+ requests + = body
144+ toServer.writeFrame(answer(request.id))
145+ }
146+ }
147+ }
148+ }
149+ supervisor.start()
150+
151+ runBlocking { server.dispatchTask(bundle, comm) }
152+ supervisor.join()
153+ comm.close()
154+ return requests to checkNotNull(terminal)
155+ }
156+
157+ private fun deleteXComRequest (
158+ key : String ,
159+ mapIndex : Int ,
160+ ) = mapOf (
161+ " type" to " DeleteXCom" ,
162+ " key" to key,
163+ " dag_id" to " dag1" ,
164+ " run_id" to " run1" ,
165+ " task_id" to " extract" ,
166+ " map_index" to mapIndex.toLong(),
167+ )
168+
63169 @Test
64170 @DisplayName(" Should run the task and report the result as a normal awaited request" )
65171 @Timeout(value = 30 , unit = TimeUnit .SECONDS )
@@ -112,6 +218,64 @@ class ServerTest {
112218 comm.close()
113219 }
114220
221+ @Test
222+ @DisplayName(" Should delete each listed XCom before the task sends anything" )
223+ @Timeout(value = 30 , unit = TimeUnit .SECONDS )
224+ fun deletesListedXComsBeforeTheTaskSendsAnything () {
225+ val cases =
226+ listOf (
227+ Triple (- 1 , listOf (" return_value" , " summary" ), - 1 ),
228+ Triple (0 , listOf (" return_value" , " summary" ), 0 ),
229+ Triple (2 , listOf (" return_value" , " summary" ), 2 ),
230+ Triple (null , listOf (" return_value" ), - 1 ),
231+ Triple (- 1 , emptyList(), - 1 ),
232+ Triple (- 1 , null , - 1 ),
233+ )
234+ for ((mapIndex, keys, deletedMapIndex) in cases) {
235+ val (requests, terminal) = serveTask(PushingTask ::class .java, mapIndex, mapOf (" xcom_keys_to_clear" to keys))
236+
237+ Assertions .assertEquals(keys.orEmpty().map { deleteXComRequest(it, deletedMapIndex) }, requests.dropLast(1 )) {
238+ " unexpected requests for map_index $mapIndex and keys $keys : $requests "
239+ }
240+ Assertions .assertEquals(" SetXCom" , requests.last()[" type" ])
241+ Assertions .assertEquals(" SucceedTask" , terminal[" type" ])
242+ }
243+ }
244+
245+ @Test
246+ @DisplayName(" Should fail without running the task when an XCom cannot be deleted" )
247+ @Timeout(value = 30 , unit = TimeUnit .SECONDS )
248+ fun failsWithoutRunningTheTaskWhenAnXComCannotBeDeleted () {
249+ for (shouldRetry in listOf (true , false )) {
250+ val (requests, terminal) =
251+ serveTask(
252+ PushingTask ::class .java,
253+ - 1 ,
254+ mapOf (" xcom_keys_to_clear" to listOf (" return_value" , " summary" ), " should_retry" to shouldRetry),
255+ ::errorResponseFrame,
256+ )
257+
258+ Assertions .assertEquals(listOf (deleteXComRequest(" return_value" , - 1 )), requests) {
259+ " the runtime should stop at the first XCom it cannot delete and not run the task"
260+ }
261+ if (shouldRetry) {
262+ Assertions .assertEquals(" RetryTask" , terminal[" type" ])
263+ } else {
264+ Assertions .assertEquals(" TaskState" , terminal[" type" ])
265+ Assertions .assertEquals(" failed" , terminal[" state" ])
266+ }
267+ }
268+ }
269+
270+ class PushingTask : Task {
271+ override fun execute (
272+ context : Context ,
273+ client : Client ,
274+ ) {
275+ client.setXCom(value = " rows" )
276+ }
277+ }
278+
115279 private companion object {
116280 // [2, msg, null] with msg coming from
117281 // https://github.com/astronomer/airflow/blob/f39c8da8/task-sdk/tests/task_sdk/execution_time/test_comms.py#L73-L108
0 commit comments