Skip to content

Commit

Permalink
fix: migrate to ReaderMock/OutputWriter
Browse files Browse the repository at this point in the history
  • Loading branch information
@namwoam
namwoam committed Sep 13, 2024
1 parent 49f7641 commit 6a3d094
Show file tree
Hide file tree
Showing 3 changed files with 128 additions and 47 deletions.
24 changes: 20 additions & 4 deletions ai/ai21labs/v0/component_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ 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 @@ -120,8 +121,16 @@ func TestComponent_Connection(t *testing.T) {
})
c.Assert(err, qt.IsNil)
pbIn := new(structpb.Struct)
_, err = exec.Execute(ctx, []*structpb.Struct{pbIn})
c.Check(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) {
return nil
})

err = exec.Execute(ctx, ir, ow)
c.Assert(err, qt.IsNil)

})

Expand All @@ -142,8 +151,15 @@ func TestComponent_Connection(t *testing.T) {
})
c.Assert(err, qt.IsNil)
pbIn := new(structpb.Struct)
_, err = exec.Execute(ctx, []*structpb.Struct{pbIn})
c.Check(err, qt.IsNotNil)
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) {
return nil
})

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

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

func (e *execution) Execute(_ context.Context, inputs []*structpb.Struct) ([]*structpb.Struct, error) {
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))

Check failure on line 130 in ai/ai21labs/v0/main.go

View workflow job for this annotation

GitHub Actions / lint 

invalid argument: inputs (variable of type *structpb.Struct) for len

Check failure on line 130 in ai/ai21labs/v0/main.go

View workflow job for this annotation

GitHub Actions / lint 

invalid argument: inputs (variable of type *structpb.Struct) for len

Check failure on line 130 in ai/ai21labs/v0/main.go

View workflow job for this annotation

GitHub Actions / codecov 

invalid argument: inputs (variable of type *structpb.Struct) for len

// The execution takes a array of inputs and returns an array of outputs. The execution is done sequentially.
for i, input := range inputs {

Check failure on line 133 in ai/ai21labs/v0/main.go

View workflow job for this annotation

GitHub Actions / lint 

cannot range over inputs (variable of type *structpb.Struct)

Check failure on line 133 in ai/ai21labs/v0/main.go

View workflow job for this annotation

GitHub Actions / lint 

cannot range over inputs (variable of type *structpb.Struct)

Check failure on line 133 in ai/ai21labs/v0/main.go

View workflow job for this annotation

GitHub Actions / codecov 

cannot range over inputs (variable of type *structpb.Struct)
output, err := e.execute(input)
if err != nil {
return nil, err
return err
}

outputs[i] = output
}

return outputs, nil
return out.Write(ctx, outputs)

Check failure on line 142 in ai/ai21labs/v0/main.go

View workflow job for this annotation

GitHub Actions / lint 

cannot use outputs (variable of type []*structpb.Struct) as *structpb.Struct value in argument to out.Write) (typecheck)

Check failure on line 142 in ai/ai21labs/v0/main.go

View workflow job for this annotation

GitHub Actions / lint 

cannot use outputs (variable of type []*structpb.Struct) as *structpb.Struct value in argument to out.Write

Check failure on line 142 in ai/ai21labs/v0/main.go

View workflow job for this annotation

GitHub Actions / codecov 

cannot use outputs (variable of type []*structpb.Struct) as *structpb.Struct value in argument to out.Write
}

// This function is not used in the codebase. It is only used in the tests.
Expand Down
141 changes: 101 additions & 40 deletions ai/ai21labs/v0/tasks_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ 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 @@ -184,12 +185,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
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) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
return nil
})

wantJSON, err := json.Marshal(tc.wantResp)
err = e.Execute(ctx, ir, ow)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
})

c.Run("ok - task embedding", func(c *qt.C) {
Expand Down Expand Up @@ -219,12 +226,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
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) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
return nil
})

wantJSON, err := json.Marshal(tc.wantResp)
err = e.Execute(ctx, ir, ow)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
})

c.Run("ok - task contextual answers without answers", func(c *qt.C) {
Expand Down Expand Up @@ -253,12 +266,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
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) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
return nil
})

wantJSON, err := json.Marshal(tc.wantResp)
err = e.Execute(ctx, ir, ow)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())

})

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

got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
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) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
return nil
})

wantJSON, err := json.Marshal(tc.wantResp)
err = e.Execute(ctx, ir, ow)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
})

c.Run("ok - task text summarization", func(c *qt.C) {
Expand All @@ -320,12 +345,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
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) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
return nil
})

wantJSON, err := json.Marshal(tc.wantResp)
err = e.Execute(ctx, ir, ow)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
})

c.Run("ok - task text summarization by segment", func(c *qt.C) {
Expand Down Expand Up @@ -358,12 +389,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
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) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
return nil
})

wantJSON, err := json.Marshal(tc.wantResp)
err = e.Execute(ctx, ir, ow)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
})

c.Run("ok - task text paraphrasing", func(c *qt.C) {
Expand All @@ -389,12 +426,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
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) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
return nil
})

wantJSON, err := json.Marshal(tc.wantResp)
err = e.Execute(ctx, ir, ow)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
})

c.Run("ok - task grammar check", func(c *qt.C) {
Expand Down Expand Up @@ -429,12 +472,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
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) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
return nil
})

wantJSON, err := json.Marshal(tc.wantResp)
err = e.Execute(ctx, ir, ow)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
})

c.Run("ok - task text improvement", func(c *qt.C) {
Expand Down Expand Up @@ -468,12 +517,18 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
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) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
return nil
})

wantJSON, err := json.Marshal(tc.wantResp)
err = e.Execute(ctx, ir, ow)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
})

c.Run("ok - task text segmentation", func(c *qt.C) {
Expand Down Expand Up @@ -505,11 +560,17 @@ func TestTasks(t *testing.T) {
pbIn, err := base.ConvertToStructpb(tc.input)
c.Assert(err, qt.IsNil)

got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
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) {
wantJSON, err := json.Marshal(tc.wantResp)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
return nil
})

wantJSON, err := json.Marshal(tc.wantResp)
err = e.Execute(ctx, ir, ow)
c.Assert(err, qt.IsNil)
c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
})
}

0 comments on commit 6a3d094

Please sign in to comment.