2020package org .apache .airflow .e2e ;
2121
2222import java .util .List ;
23+ import java .util .Map ;
24+ import java .util .Objects ;
2325import org .apache .airflow .sdk .*;
2426import org .jetbrains .annotations .NotNull ;
2527
2628/**
2729 * Bundle for the runner-behaviour E2E tests: deliberately broken task classes that exercise
28- * instantiation failures, and a task that round-trips Airflow Variables through the supervisor.
30+ * instantiation failures, and tasks that round-trip Airflow Variables and asset state through the
31+ * supervisor.
2932 */
3033public class TestBundleBuilder implements BundleBuilder {
3134 public static class MissingNoArgConstructor implements Task {
@@ -60,6 +63,46 @@ public void execute(@NotNull Context context, Client client) {
6063 }
6164 }
6265
66+ /**
67+ * Calls every asset state store method on one asset, through a lookup by name and a lookup by
68+ * URI. Each step changes the state through one lookup and checks the result through the other,
69+ * so the task fails unless both lookups reach the same asset. The E2E test then reads the
70+ * remaining keys through the REST API.
71+ */
72+ public static class UseAssetStateStore implements Task {
73+ public void execute (@ NotNull Context context , Client client ) {
74+ var byName = client .assetStateStore .byName ("java_e2e_orders" );
75+ var byUri = client .assetStateStore .byUri ("x-java-e2e://orders" );
76+
77+ byName .set ("scratch" , "first" );
78+ expect ("first" , byUri .get ("scratch" ), "set() through the name" );
79+ byUri .clear ();
80+ expect (null , byName .get ("scratch" ), "clear() through the URI" );
81+
82+ byUri .set ("scratch" , "second" );
83+ expect ("second" , byName .get ("scratch" ), "set() through the URI" );
84+ byName .clear ();
85+ expect (null , byUri .get ("scratch" ), "clear() through the name" );
86+
87+ byName .set ("deleted" , "temporary" );
88+ byUri .delete ("deleted" );
89+ expect (null , byName .get ("deleted" ), "delete() through the URI" );
90+ byUri .set ("deleted" , "temporary" );
91+ byName .delete ("deleted" );
92+ expect (null , byUri .get ("deleted" ), "delete() through the name" );
93+
94+ var summary = Map .of ("run_id" , context .dagRun .runId , "rows" , 3L );
95+ byName .set ("summary" , summary );
96+ expect (summary , byUri .get ("summary" ), "set() through the name, with a map value" );
97+ }
98+
99+ private static void expect (Object expected , Object actual , String step ) {
100+ if (!Objects .equals (expected , actual )) {
101+ throw new IllegalStateException (step + ": expected " + expected + " but read " + actual );
102+ }
103+ }
104+ }
105+
63106 @ NotNull
64107 @ Override
65108 public Iterable <DagDef > getDags () {
@@ -68,7 +111,9 @@ public Iterable<DagDef> getDags() {
68111 uninstantiable .addTask ("non_static_inner" , NonStaticInner .class );
69112 var variableWrite = new DagDef ("java_variable_write" );
70113 variableWrite .addTask ("write_and_delete" , WriteAndDeleteVariable .class );
71- return List .of (uninstantiable , variableWrite );
114+ var assetStateStore = new DagDef ("java_asset_state_store" );
115+ assetStateStore .addTask ("use_asset_state_store" , UseAssetStateStore .class );
116+ return List .of (uninstantiable , variableWrite , assetStateStore );
72117 }
73118
74119 public static void main (String [] args ) {
0 commit comments