55import io .cucumber .messages .types .Examples ;
66import io .cucumber .messages .types .Feature ;
77import io .cucumber .messages .types .GherkinDocument ;
8+ import io .cucumber .messages .types .Hook ;
9+ import io .cucumber .messages .types .Location ;
810import io .cucumber .messages .types .Pickle ;
911import io .cucumber .messages .types .PickleStep ;
1012import io .cucumber .messages .types .Rule ;
2022import io .cucumber .messages .types .TestStepFinished ;
2123import io .cucumber .messages .types .TestStepResult ;
2224import io .cucumber .messages .types .TestStepResultStatus ;
25+ import io .cucumber .messages .types .TestStepStarted ;
2326import io .cucumber .messages .types .Timestamp ;
2427import io .cucumber .query .LineageReducerStrategy .Descending ;
2528
6467public final class Query {
6568 private final Comparator <TestStepResult > testStepResultComparator = nullsFirst (comparing (o -> o .getStatus ().ordinal ()));
6669 private final Deque <TestCaseStarted > testCaseStarted = new ConcurrentLinkedDeque <>();
70+ private final Map <String , TestCaseStarted > testCaseStartedById = new ConcurrentHashMap <>();
6771 private final Map <String , TestCaseFinished > testCaseFinishedByTestCaseStartedId = new ConcurrentHashMap <>();
6872 private final Map <String , List <TestStepFinished >> testStepsFinishedByTestCaseStartedId = new ConcurrentHashMap <>();
6973 private final Map <String , Pickle > pickleById = new ConcurrentHashMap <>();
@@ -72,6 +76,7 @@ public final class Query {
7276 private final Map <String , TestStep > testStepById = new ConcurrentHashMap <>();
7377 private final Map <String , PickleStep > pickleStepById = new ConcurrentHashMap <>();
7478 private final Map <Object , Lineage > lineageById = new ConcurrentHashMap <>();
79+ private final Map <String , Hook > hookById = new ConcurrentHashMap <>();
7580 private TestRunStarted testRunStarted ;
7681 private TestRunFinished testRunFinished ;
7782
@@ -137,6 +142,11 @@ public Optional<Feature> findFeatureBy(TestCaseStarted testCaseStarted) {
137142 return findLineageBy (testCaseStarted ).flatMap (Lineage ::feature );
138143 }
139144
145+ public Optional <Hook > findHookBy (TestStep testStep ) {
146+ requireNonNull (testStep );
147+ return testStep .getHookId ().map (hookById ::get );
148+ }
149+
140150 public Optional <TestStepResult > findMostSevereTestStepResulBy (TestCaseStarted testCaseStarted ) {
141151 requireNonNull (testCaseStarted );
142152 return findTestStepsFinishedBy (testCaseStarted )
@@ -193,12 +203,16 @@ public String findNameOf(TableRow element, NamingStrategy namingStrategy) {
193203 .orElseThrow (createElementWasNotPartOfThisQueryObject ());
194204 }
195205
196- public String findNameOf (Pickle element , NamingStrategy namingStrategy ) {
206+ public Optional < String > findNameOf (Pickle element , NamingStrategy namingStrategy ) {
197207 requireNonNull (element );
198208 requireNonNull (namingStrategy );
199209 return findLineageBy (element )
200- .map (lineage -> namingStrategy .reduce (lineage , element ))
201- .orElseThrow (createElementWasNotPartOfThisQueryObject ());
210+ .map (lineage -> namingStrategy .reduce (lineage , element ));
211+ }
212+
213+ public Optional <Location > findLocationOf (Pickle element ) {
214+ requireNonNull (element );
215+ return reduceLinageOf (element , LocationOf ::new );
202216 }
203217
204218 private static Supplier <IllegalArgumentException > createElementWasNotPartOfThisQueryObject () {
@@ -254,7 +268,7 @@ <T> Optional<T> reduceLinageOf(TableRow element, Supplier<LineageReducer<T>> red
254268 .map (strategy ::reduce );
255269 }
256270
257- <T > Optional <T > reduceLinageOf (Pickle element , Supplier <LineageReducer <T >> reducerSupplier ) {
271+ public <T > Optional <T > reduceLinageOf (Pickle element , Supplier <LineageReducer <T >> reducerSupplier ) {
258272 requireNonNull (element );
259273 requireNonNull (reducerSupplier );
260274 Descending <T > strategy = new Descending <>(reducerSupplier );
@@ -269,8 +283,15 @@ public Optional<Pickle> findPickleBy(TestCaseStarted testCaseStarted) {
269283 .map (pickleById ::get );
270284 }
271285
286+ public Optional <Pickle > findPickleBy (TestStepStarted testStepStarted ) {
287+ requireNonNull (testStepStarted );
288+ return findTestCaseBy (testStepStarted )
289+ .map (TestCase ::getPickleId )
290+ .map (pickleById ::get );
291+ }
292+
272293 public Optional <PickleStep > findPickleStepBy (TestStep testStep ) {
273- requireNonNull (testCaseStarted );
294+ requireNonNull (testStep );
274295 return testStep .getPickleStepId ()
275296 .map (pickleStepById ::get );
276297 }
@@ -281,11 +302,27 @@ public Optional<Step> findStepBy(PickleStep pickleStep) {
281302 return ofNullable (stepById .get (stepId ));
282303 }
283304
305+ public Optional <Step > findStepBy (TestStep testStep ) {
306+ requireNonNull (testStep );
307+ return findPickleStepBy (testStep )
308+ .flatMap (this ::findStepBy );
309+ }
310+
284311 public Optional <TestCase > findTestCaseBy (TestCaseStarted testCaseStarted ) {
285312 requireNonNull (testCaseStarted );
286313 return ofNullable (testCaseById .get (testCaseStarted .getTestCaseId ()));
287314 }
288315
316+ public Optional <TestCase > findTestCaseBy (TestStepStarted testCaseStarted ) {
317+ return findTestCaseStartedBy (testCaseStarted )
318+ .flatMap (this ::findTestCaseBy );
319+ }
320+
321+ public Optional <TestCaseStarted > findTestCaseStartedBy (TestStepStarted testStepStarted ) {
322+ requireNonNull (testStepStarted );
323+ return ofNullable (testCaseStartedById .get (testStepStarted .getTestCaseStartedId ()));
324+ }
325+
289326 public Optional <Duration > findTestCaseDurationBy (TestCaseStarted testCaseStarted ) {
290327 requireNonNull (testCaseStarted );
291328 Timestamp started = testCaseStarted .getTimestamp ();
@@ -321,6 +358,11 @@ public Optional<TestRunStarted> findTestRunStarted() {
321358 return ofNullable (testRunStarted );
322359 }
323360
361+ public Optional <TestStep > findTestStepBy (TestStepStarted testStepStarted ) {
362+ requireNonNull (testStepStarted );
363+ return ofNullable (testStepById .get (testStepStarted .getTestStepId ()));
364+ }
365+
324366 public Optional <TestStep > findTestStepBy (TestStepFinished testStepFinished ) {
325367 requireNonNull (testStepFinished );
326368 return ofNullable (testStepById .get (testStepFinished .getTestStepId ()));
@@ -351,6 +393,7 @@ public void update(Envelope envelope) {
351393 envelope .getGherkinDocument ().ifPresent (this ::updateGherkinDocument );
352394 envelope .getPickle ().ifPresent (this ::updatePickle );
353395 envelope .getTestCase ().ifPresent (this ::updateTestCase );
396+ envelope .getHook ().ifPresent (this ::updateHook );
354397 }
355398
356399 private Optional <Lineage > findLineageBy (GherkinDocument element ) {
@@ -397,6 +440,7 @@ private Optional<Lineage> findLineageBy(TestCaseStarted testCaseStarted) {
397440
398441 private void updateTestCaseStarted (TestCaseStarted testCaseStarted ) {
399442 this .testCaseStarted .add (testCaseStarted );
443+ this .testCaseStartedById .put (testCaseStarted .getId (), testCaseStarted );
400444 }
401445
402446 private void updateTestCase (TestCase event ) {
@@ -450,6 +494,10 @@ private void updateSteps(List<Step> steps) {
450494 steps .forEach (step -> stepById .put (step .getId (), step ));
451495 }
452496
497+ private void updateHook (Hook event ) {
498+ this .hookById .put (event .getId (), event );
499+ }
500+
453501 private <K , E > BiFunction <K , List <E >, List <E >> updateList (E element ) {
454502 return (key , existing ) -> {
455503 if (existing != null ) {
0 commit comments