Skip to content

Commit

Permalink
refactor: migrate to job and error handler
Browse files Browse the repository at this point in the history
  • Loading branch information
@namwoam
namwoam committed Sep 13, 2024
1 parent 7fc8fe2 commit 467606d
Show file tree
Hide file tree
Showing 3 changed files with 87 additions and 93 deletions.
29 changes: 15 additions & 14 deletions ai/ai21labs/v0/component_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@ import (

qt "github.com/frankban/quicktest"
"github.com/instill-ai/component/base"
"github.com/instill-ai/component/internal/mock"
"github.com/instill-ai/component/internal/util/httpclient"
"github.com/instill-ai/x/errmsg"
"go.uber.org/zap"
Expand Down Expand Up @@ -114,22 +113,23 @@ func TestComponent_Connection(t *testing.T) {
})
c.Assert(err, qt.IsNil)

exec, err := cmp.CreateExecution(base.ComponentExecution{
e, err := cmp.CreateExecution(base.ComponentExecution{
Component: cmp,
Task: "TASK_TEXT_GENERATION_CHAT",
Setup: setup,
})
c.Assert(err, qt.IsNil)
pbIn := new(structpb.Struct)

ir := mock.NewInputReaderMock(c)
ow := mock.NewOutputWriterMock(c)
ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
ir, ow, eh, job := base.GenerateMockJob(c)
ir.ReadMock.Return(pbIn, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, output *structpb.Struct) (err error) {
return nil
})

err = exec.Execute(ctx, ir, ow)
eh.ErrorMock.Optional()

err = e.Execute(ctx, []*base.Job{job})
c.Assert(err, qt.IsNil)

})
Expand All @@ -144,22 +144,23 @@ func TestComponent_Connection(t *testing.T) {
})
c.Assert(err, qt.IsNil)

exec, err := cmp.CreateExecution(base.ComponentExecution{
e, err := cmp.CreateExecution(base.ComponentExecution{
Component: cmp,
Task: "TASK_TEXT_GENERATION_CHAT",
Setup: setup,
})
c.Assert(err, qt.IsNil)
pbIn := new(structpb.Struct)
ir := mock.NewInputReaderMock(c)
ow := mock.NewOutputWriterMock(c)
ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
ir, ow, eh, job := base.GenerateMockJob(c)
ir.ReadMock.Return(pbIn, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, output *structpb.Struct) (err error) {
return nil
})

err = exec.Execute(ctx, ir, ow)
c.Assert(err, qt.IsNotNil)
eh.ErrorMock.Optional()

err = e.Execute(ctx, []*base.Job{job})
c.Assert(err, qt.IsNil)

want := "AI21labs responded with a 401 status code. Incorrect API key provided."
c.Check(errmsg.Message(err), qt.Equals, want)
Expand Down
20 changes: 2 additions & 18 deletions ai/ai21labs/v0/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,24 +122,8 @@ func (e *execution) UsesInstillCredentials() bool {
return e.usesInstillCredentials
}

func (e *execution) Execute(ctx context.Context, in base.InputReader, out base.OutputWriter) error {
inputs, err := in.Read(ctx)
if err != nil {
return err
}
outputs := make([]*structpb.Struct, len(inputs))

// The execution takes a array of inputs and returns an array of outputs. The execution is done sequentially.
for i, input := range inputs {
output, err := e.execute(input)
if err != nil {
return err
}

outputs[i] = output
}

return out.Write(ctx, outputs)
func (e *execution) Execute(ctx context.Context, jobs []*base.Job) error {
return base.SequentialExecutor(ctx, jobs, e.execute)
}

// This function is not used in the codebase. It is only used in the tests.
Expand Down
131 changes: 70 additions & 61 deletions ai/ai21labs/v0/tasks_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@ import (
qt "github.com/frankban/quicktest"
"github.com/instill-ai/component/ai"
"github.com/instill-ai/component/base"
"github.com/instill-ai/component/internal/mock"
"go.uber.org/zap"
"google.golang.org/protobuf/types/known/structpb"
)
Expand Down Expand Up @@ -185,17 +184,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

ir := mock.NewInputReaderMock(c)
ow := mock.NewOutputWriterMock(c)
ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
ir, ow, eh, job := base.GenerateMockJob(c)
ir.ReadMock.Return(pbIn, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, output *structpb.Struct) (err error) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
c.Check(wantJSON, qt.JSONEquals, output.AsMap())
return nil
})

err = e.Execute(ctx, ir, ow)
eh.ErrorMock.Optional()

err = e.Execute(ctx, []*base.Job{job})
c.Assert(err, qt.IsNil)
})

Expand Down Expand Up @@ -226,17 +226,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

ir := mock.NewInputReaderMock(c)
ow := mock.NewOutputWriterMock(c)
ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
ir, ow, eh, job := base.GenerateMockJob(c)
ir.ReadMock.Return(pbIn, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, output *structpb.Struct) (err error) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
c.Check(wantJSON, qt.JSONEquals, output.AsMap())
return nil
})

err = e.Execute(ctx, ir, ow)
eh.ErrorMock.Optional()

err = e.Execute(ctx, []*base.Job{job})
c.Assert(err, qt.IsNil)
})

Expand Down Expand Up @@ -266,17 +267,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

ir := mock.NewInputReaderMock(c)
ow := mock.NewOutputWriterMock(c)
ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
ir, ow, eh, job := base.GenerateMockJob(c)
ir.ReadMock.Return(pbIn, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, output *structpb.Struct) (err error) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
c.Check(wantJSON, qt.JSONEquals, output.AsMap())
return nil
})

err = e.Execute(ctx, ir, ow)
eh.ErrorMock.Optional()

err = e.Execute(ctx, []*base.Job{job})
c.Assert(err, qt.IsNil)

})
Expand Down Expand Up @@ -307,17 +309,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

ir := mock.NewInputReaderMock(c)
ow := mock.NewOutputWriterMock(c)
ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
ir, ow, eh, job := base.GenerateMockJob(c)
ir.ReadMock.Return(pbIn, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, output *structpb.Struct) (err error) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
c.Check(wantJSON, qt.JSONEquals, output.AsMap())
return nil
})

err = e.Execute(ctx, ir, ow)
eh.ErrorMock.Optional()

err = e.Execute(ctx, []*base.Job{job})
c.Assert(err, qt.IsNil)
})

Expand Down Expand Up @@ -345,17 +348,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

ir := mock.NewInputReaderMock(c)
ow := mock.NewOutputWriterMock(c)
ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
ir, ow, eh, job := base.GenerateMockJob(c)
ir.ReadMock.Return(pbIn, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, output *structpb.Struct) (err error) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
c.Check(wantJSON, qt.JSONEquals, output.AsMap())
return nil
})

err = e.Execute(ctx, ir, ow)
eh.ErrorMock.Optional()

err = e.Execute(ctx, []*base.Job{job})
c.Assert(err, qt.IsNil)
})

Expand Down Expand Up @@ -389,17 +393,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

ir := mock.NewInputReaderMock(c)
ow := mock.NewOutputWriterMock(c)
ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
ir, ow, eh, job := base.GenerateMockJob(c)
ir.ReadMock.Return(pbIn, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, output *structpb.Struct) (err error) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
c.Check(wantJSON, qt.JSONEquals, output.AsMap())
return nil
})

err = e.Execute(ctx, ir, ow)
eh.ErrorMock.Optional()

err = e.Execute(ctx, []*base.Job{job})
c.Assert(err, qt.IsNil)
})

Expand All @@ -426,17 +431,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

ir := mock.NewInputReaderMock(c)
ow := mock.NewOutputWriterMock(c)
ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
ir, ow, eh, job := base.GenerateMockJob(c)
ir.ReadMock.Return(pbIn, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, output *structpb.Struct) (err error) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
c.Check(wantJSON, qt.JSONEquals, output.AsMap())
return nil
})

err = e.Execute(ctx, ir, ow)
eh.ErrorMock.Optional()

err = e.Execute(ctx, []*base.Job{job})
c.Assert(err, qt.IsNil)
})

Expand Down Expand Up @@ -472,17 +478,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

ir := mock.NewInputReaderMock(c)
ow := mock.NewOutputWriterMock(c)
ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
ir, ow, eh, job := base.GenerateMockJob(c)
ir.ReadMock.Return(pbIn, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, output *structpb.Struct) (err error) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
c.Check(wantJSON, qt.JSONEquals, output.AsMap())
return nil
})

err = e.Execute(ctx, ir, ow)
eh.ErrorMock.Optional()

err = e.Execute(ctx, []*base.Job{job})
c.Assert(err, qt.IsNil)
})

Expand Down Expand Up @@ -517,17 +524,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

ir := mock.NewInputReaderMock(c)
ow := mock.NewOutputWriterMock(c)
ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
ir, ow, eh, job := base.GenerateMockJob(c)
ir.ReadMock.Return(pbIn, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, output *structpb.Struct) (err error) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
c.Check(wantJSON, qt.JSONEquals, output.AsMap())
return nil
})

err = e.Execute(ctx, ir, ow)
eh.ErrorMock.Optional()

err = e.Execute(ctx, []*base.Job{job})
c.Assert(err, qt.IsNil)
})

Expand Down Expand Up @@ -560,17 +568,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

ir := mock.NewInputReaderMock(c)
ow := mock.NewOutputWriterMock(c)
ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
ir, ow, eh, job := base.GenerateMockJob(c)
ir.ReadMock.Return(pbIn, nil)
ow.WriteMock.Optional().Set(func(ctx context.Context, output *structpb.Struct) (err error) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
c.Check(wantJSON, qt.JSONEquals, output.AsMap())
return nil
})

err = e.Execute(ctx, ir, ow)
eh.ErrorMock.Optional()

err = e.Execute(ctx, []*base.Job{job})
c.Assert(err, qt.IsNil)
})
}

0 comments on commit 467606d

Please sign in to comment.