22
22
import static com .google .cloud .datastore .telemetry .TraceUtil .SPAN_NAME_COMMIT ;
23
23
import static com .google .cloud .datastore .telemetry .TraceUtil .SPAN_NAME_LOOKUP ;
24
24
import static com .google .cloud .datastore .telemetry .TraceUtil .SPAN_NAME_RESERVE_IDS ;
25
+ import static com .google .cloud .datastore .telemetry .TraceUtil .SPAN_NAME_ROLLBACK ;
25
26
import static com .google .cloud .datastore .telemetry .TraceUtil .SPAN_NAME_RUN_AGGREGATION_QUERY ;
26
27
import static com .google .cloud .datastore .telemetry .TraceUtil .SPAN_NAME_RUN_QUERY ;
27
28
import static com .google .cloud .datastore .telemetry .TraceUtil .SPAN_NAME_TRANSACTION_COMMIT ;
@@ -374,6 +375,7 @@ public void after() throws Exception {
374
375
tracer = null ;
375
376
retrievedTrace = null ;
376
377
customSpanContext = null ;
378
+ openTelemetrySdk = null ;
377
379
}
378
380
379
381
@ AfterClass
@@ -793,7 +795,7 @@ public void runAggregationQueryTraceTest() throws Exception {
793
795
}
794
796
795
797
@ Test
796
- public void transactionalLookupTest () throws Exception {
798
+ public void newTransactionReadTest () throws Exception {
797
799
assertNotNull (customSpanContext );
798
800
799
801
Span rootSpan = getNewRootSpanWithContext ();
@@ -817,7 +819,7 @@ public void transactionalLookupTest() throws Exception {
817
819
}
818
820
819
821
@ Test
820
- public void transactionQueryTest () throws Exception {
822
+ public void newTransactionQueryTest () throws Exception {
821
823
// Set up
822
824
Entity entity1 = Entity .newBuilder (KEY1 ).set ("test_field" , "test_value1" ).build ();
823
825
Entity entity2 = Entity .newBuilder (KEY2 ).set ("test_field" , "test_value2" ).build ();
@@ -856,6 +858,116 @@ public void transactionQueryTest() throws Exception {
856
858
Collections .singletonList (SPAN_NAME_TRANSACTION_COMMIT )));
857
859
}
858
860
861
+ @ Test
862
+ public void newTransactionReadWriteTraceTest () throws Exception {
863
+ // Set up
864
+ Entity entity1 = Entity .newBuilder (KEY1 ).set ("pepper_type" , "jalapeno" ).build ();
865
+ Entity entity2 = Entity .newBuilder (KEY2 ).set ("pepper_type" , "habanero" ).build ();
866
+ List <Entity > entityList = new ArrayList <>();
867
+ entityList .add (entity1 );
868
+ entityList .add (entity2 );
869
+
870
+ List <Entity > response = datastore .add (entity1 , entity2 );
871
+ assertEquals (entityList , response );
872
+
873
+ String simplified_spice_level = "not_spicy" ;
874
+ Entity entity1update =
875
+ Entity .newBuilder (entity1 ).set ("spice_level" , simplified_spice_level ).build ();
876
+
877
+ assertNotNull (customSpanContext );
878
+
879
+ // Test
880
+ Span rootSpan = getNewRootSpanWithContext ();
881
+ try (Scope ignored = rootSpan .makeCurrent ()) {
882
+ Transaction transaction = datastore .newTransaction ();
883
+ entity1 = transaction .get (KEY1 );
884
+ switch (entity1 .getString ("pepper_type" )) {
885
+ case "jalapeno" :
886
+ simplified_spice_level = "mild" ;
887
+ break ;
888
+
889
+ case "habanero" :
890
+ simplified_spice_level = "hot" ;
891
+ break ;
892
+ }
893
+ transaction .update (entity1update );
894
+ transaction .delete (KEY2 );
895
+ transaction .commit ();
896
+ assertFalse (transaction .isActive ());
897
+ } finally {
898
+ rootSpan .end ();
899
+ }
900
+
901
+ waitForTracesToComplete ();
902
+
903
+ List <Entity > list = datastore .fetch (KEY1 , KEY2 );
904
+ assertEquals (list .get (0 ), entity1update );
905
+ assertNull (list .get (1 ));
906
+
907
+ fetchAndValidateTrace (
908
+ customSpanContext .getTraceId (),
909
+ /*numExpectedSpans=*/ 3 ,
910
+ Arrays .asList (
911
+ Collections .singletonList (SPAN_NAME_BEGIN_TRANSACTION ),
912
+ Collections .singletonList (SPAN_NAME_TRANSACTION_LOOKUP ),
913
+ Collections .singletonList (SPAN_NAME_TRANSACTION_COMMIT )));
914
+ }
915
+
916
+ @ Test
917
+ public void newTransactionRollbackTest () throws Exception {
918
+ // Set up
919
+ Entity entity1 = Entity .newBuilder (KEY1 ).set ("pepper_type" , "jalapeno" ).build ();
920
+ Entity entity2 = Entity .newBuilder (KEY2 ).set ("pepper_type" , "habanero" ).build ();
921
+ List <Entity > entityList = new ArrayList <>();
922
+ entityList .add (entity1 );
923
+ entityList .add (entity2 );
924
+
925
+ List <Entity > response = datastore .add (entity1 , entity2 );
926
+ assertEquals (entityList , response );
927
+
928
+ String simplified_spice_level = "not_spicy" ;
929
+ Entity entity1update =
930
+ Entity .newBuilder (entity1 ).set ("spice_level" , simplified_spice_level ).build ();
931
+
932
+ assertNotNull (customSpanContext );
933
+
934
+ // Test
935
+ Span rootSpan = getNewRootSpanWithContext ();
936
+ try (Scope ignored = rootSpan .makeCurrent ()) {
937
+ Transaction transaction = datastore .newTransaction ();
938
+ entity1 = transaction .get (KEY1 );
939
+ switch (entity1 .getString ("pepper_type" )) {
940
+ case "jalapeno" :
941
+ simplified_spice_level = "mild" ;
942
+ break ;
943
+
944
+ case "habanero" :
945
+ simplified_spice_level = "hot" ;
946
+ break ;
947
+ }
948
+ transaction .update (entity1update );
949
+ transaction .delete (KEY2 );
950
+ transaction .rollback ();
951
+ assertFalse (transaction .isActive ());
952
+ } finally {
953
+ rootSpan .end ();
954
+ }
955
+
956
+ waitForTracesToComplete ();
957
+
958
+ List <Entity > list = datastore .fetch (KEY1 , KEY2 );
959
+ assertEquals (list .get (0 ), entity1 );
960
+ assertEquals (list .get (1 ), entity2 );
961
+
962
+ fetchAndValidateTrace (
963
+ customSpanContext .getTraceId (),
964
+ /*numExpectedSpans=*/ 3 ,
965
+ Arrays .asList (
966
+ Collections .singletonList (SPAN_NAME_BEGIN_TRANSACTION ),
967
+ Collections .singletonList (SPAN_NAME_TRANSACTION_LOOKUP ),
968
+ Collections .singletonList (SPAN_NAME_ROLLBACK )));
969
+ }
970
+
859
971
@ Test
860
972
public void runInTransactionQueryTest () throws Exception {
861
973
// Set up
@@ -897,10 +1009,4 @@ public void runInTransactionQueryTest() throws Exception {
897
1009
Arrays .asList (SPAN_NAME_TRANSACTION_RUN , SPAN_NAME_RUN_QUERY ),
898
1010
Arrays .asList (SPAN_NAME_TRANSACTION_RUN , SPAN_NAME_TRANSACTION_COMMIT )));
899
1011
}
900
-
901
- @ Test
902
- public void transactionRunQueryTest () throws Exception {}
903
-
904
- @ Test
905
- public void readWriteTransactionTraceTest () throws Exception {}
906
1012
}
0 commit comments