From 2d703e7f6596071aef1a752b6d9c4ba470e644e8 Mon Sep 17 00:00:00 2001
From: gabriel ruttner <gabriel.ruttner@gmail.com>
Date: Wed, 18 Sep 2024 15:36:36 -0400
Subject: [PATCH] fix: tests

---
 examples/cancellation/run.go |  8 ++++++--
 examples/timeout/run.go      | 14 ++++++++++----
 examples/webhook/run.go      | 14 ++++++++++----
 3 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/examples/cancellation/run.go b/examples/cancellation/run.go
index f5995ccd3..b134f8319 100644
--- a/examples/cancellation/run.go
+++ b/examples/cancellation/run.go
@@ -102,8 +102,12 @@ func run(events chan<- string) (func() error, error) {
 
 		for _, stepRun := range stepRuns {
 			stepRunID := stepRun.ID
-			log.Printf("cancelling step run id: %s", stepRunID)
-			res, err := c.API().StepRunUpdateCancelWithResponse(context.Background(), uuid.MustParse(c.TenantId()), uuid.MustParse(stepRunID))
+			log.Printf("cancelling step run id: %d", stepRunID)
+			res, err := c.API().StepRunUpdateCancelWithResponse(
+				context.Background(),
+				uuid.MustParse(c.TenantId()),
+				fmt.Sprint(stepRunID),
+			)
 			if err != nil {
 				panic(fmt.Errorf("error cancelling step run: %w", err))
 			}
diff --git a/examples/timeout/run.go b/examples/timeout/run.go
index 5e270ad71..2451ae901 100644
--- a/examples/timeout/run.go
+++ b/examples/timeout/run.go
@@ -71,12 +71,11 @@ func run(done chan<- string, job worker.WorkflowJob) (func() error, error) {
 		).With(
 			db.Event.WorkflowRuns.Fetch().With(
 				db.WorkflowRunTriggeredBy.Parent.Fetch().With(
-					db.WorkflowRun.JobRuns.Fetch().With(
-						db.JobRun.StepRuns.Fetch(),
-					),
+					db.WorkflowRun.JobRuns.Fetch(),
 				),
 			),
 		).Exec(context.Background())
+
 		if err != nil {
 			panic(fmt.Errorf("error finding events: %w", err))
 		}
@@ -84,7 +83,14 @@ func run(done chan<- string, job worker.WorkflowJob) (func() error, error) {
 		for _, event := range events {
 			for _, workflowRun := range event.WorkflowRuns() {
 				for _, jobRuns := range workflowRun.Parent().JobRuns() {
-					for _, stepRun := range jobRuns.StepRuns() {
+					stepRuns, err := client.StepRun.FindMany(
+						db.StepRun.JobRunID.Equals(jobRuns.ID),
+					).Exec(context.Background())
+					if err != nil {
+						panic(fmt.Errorf("error finding step runs: %w", err))
+					}
+
+					for _, stepRun := range stepRuns {
 						if stepRun.Status != db.StepRunStatusFailed {
 							panic(fmt.Errorf("expected step run to be failed, got %s", stepRun.Status))
 						}
diff --git a/examples/webhook/run.go b/examples/webhook/run.go
index ff17782d7..a6c6ca485 100644
--- a/examples/webhook/run.go
+++ b/examples/webhook/run.go
@@ -94,9 +94,7 @@ func verifyStepRuns(prisma *db.PrismaClient, event string, tenantId string, jobR
 	).With(
 		db.Event.WorkflowRuns.Fetch().With(
 			db.WorkflowRunTriggeredBy.Parent.Fetch().With(
-				db.WorkflowRun.JobRuns.Fetch().With(
-					db.JobRun.StepRuns.Fetch(),
-				),
+				db.WorkflowRun.JobRuns.Fetch(),
 			),
 		),
 	).Exec(context.Background())
@@ -120,7 +118,15 @@ func verifyStepRuns(prisma *db.PrismaClient, event string, tenantId string, jobR
 				if jobRuns.Status != jobRunStatus {
 					panic(fmt.Errorf("expected job run to be %s, got %s", jobRunStatus, jobRuns.Status))
 				}
-				for _, stepRun := range jobRuns.StepRuns() {
+
+				stepRuns, err := prisma.StepRun.FindMany(
+					db.StepRun.JobRunID.Equals(jobRuns.ID),
+				).Exec(context.Background())
+				if err != nil {
+					panic(fmt.Errorf("error finding step runs: %w", err))
+				}
+
+				for _, stepRun := range stepRuns {
 					if stepRun.Status != stepRunStatus {
 						panic(fmt.Errorf("expected step run to be %s, got %s", stepRunStatus, stepRun.Status))
 					}