diff --git a/test/integration_tests/executor_test.go b/test/integration_tests/executor_test.go
index 91f89904..f7c7d5e2 100644
--- a/test/integration_tests/executor_test.go
+++ b/test/integration_tests/executor_test.go
@@ -24,7 +24,7 @@ func TestRetryNotFound(t *testing.T) {
 	assert.Error(t, err, "Retry is expected to return an error")
 
 	if swaggerErr, ok := err.(client.GenericSwaggerError); ok {
-		assert.Equal(t, 500, swaggerErr.StatusCode()) //TODO on test server update, should be a 404
+		assert.Equal(t, 404, swaggerErr.StatusCode())
 	} else {
 		assert.Fail(t, "err is not of type GenericSwaggerError")
 	}
@@ -74,7 +74,7 @@ func TestUpdateTaskByRefName(t *testing.T) {
 	err := executor.UpdateTaskByRefName("task_ref", notFoundWorkflowId, model.CompletedTask, map[string]interface{}{})
 	assert.Error(t, err, "UpdateTaskByRefName is expected to return an error")
 	if swaggerErr, ok := err.(client.GenericSwaggerError); ok {
-		assert.Equal(t, 500, swaggerErr.StatusCode()) //TODO on test server update, should be a 404
+		assert.Equal(t, 404, swaggerErr.StatusCode())
 	} else {
 		assert.Fail(t, "err is not of type GenericSwaggerError")
 	}
diff --git a/test/integration_tests/workflow_definition_test.go b/test/integration_tests/workflow_definition_test.go
index a1b48f84..d4ff6928 100644
--- a/test/integration_tests/workflow_definition_test.go
+++ b/test/integration_tests/workflow_definition_test.go
@@ -18,6 +18,7 @@ import (
 const retryLimit = 5
 
 func TestWorkflowCreation(t *testing.T) {
+	executor := testdata.WorkflowExecutor
 	workflow := testdata.NewKitchenSinkWorkflow(testdata.WorkflowExecutor)
 	err := workflow.Register(true)
 	if err != nil {
@@ -31,9 +32,25 @@ func TestWorkflowCreation(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Failed to complete the workflow, reason: %s", err)
 	}
+
 	assert.NotEmpty(t, run, "Workflow is null", run)
-	assert.Equal(t, string(model.CompletedWorkflow), run.Status)
-	assert.Equal(t, "input1", run.Input["key1"])
+
+	timeout := time.After(60 * time.Second)
+	tick := time.Tick(1 * time.Second)
+	workflowId := run.WorkflowId
+	for {
+		select {
+		case <-timeout:
+			t.Fatalf("Timed out and workflow %s didn't complete", workflowId)
+		case <-tick:
+			wf, err := executor.GetWorkflow(workflowId, false)
+			assert.NoError(t, err)
+			assert.Equal(t, model.CompletedWorkflow, wf.Status)
+			assert.Equal(t, "input1", run.Input["key1"])
+			return
+		}
+	}
+
 }
 
 func TestRemoveWorkflow(t *testing.T) {