github.com/instill-ai/component@v0.16.0-beta/pkg/connector/restapi/v0/connector_test.go (about) 1 package restapi 2 3 import ( 4 "encoding/base64" 5 "fmt" 6 "io" 7 "net/http" 8 "net/http/httptest" 9 "testing" 10 11 qt "github.com/frankban/quicktest" 12 "go.uber.org/zap" 13 "google.golang.org/protobuf/types/known/structpb" 14 15 "github.com/instill-ai/component/pkg/base" 16 "github.com/instill-ai/component/pkg/connector/util/httpclient" 17 "github.com/instill-ai/x/errmsg" 18 ) 19 20 const ( 21 errResp = `{"message": "Bad request"}` 22 okResp = `{"title": "Be the wheel"}` 23 ) 24 25 var ( 26 path = "/good-songs/10" 27 ) 28 29 func TestConnector_Execute(t *testing.T) { 30 c := qt.New(t) 31 32 logger := zap.NewNop() 33 connector := Init(logger, nil) 34 reqBody := map[string]any{ 35 "title": "Be the wheel", 36 } 37 38 c.Run("nok - unsupported task", func(c *qt.C) { 39 task := "FOOBAR" 40 41 exec, err := connector.CreateExecution(nil, cfg(noAuthType), task) 42 c.Assert(err, qt.IsNil) 43 44 pbIn := new(structpb.Struct) 45 _, err = exec.Execution.Execute([]*structpb.Struct{pbIn}) 46 c.Check(err, qt.IsNotNil) 47 48 want := "FOOBAR task is not supported." 49 c.Check(errmsg.Message(err), qt.Equals, want) 50 }) 51 52 c.Run("ok - POST, 400, basic auth", func(c *qt.C) { 53 h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 54 c.Check(r.Method, qt.Equals, http.MethodPost) 55 c.Check(r.URL.Path, qt.Matches, path) 56 57 auth := base64.StdEncoding.EncodeToString([]byte(username + ":" + password)) 58 c.Check(r.Header.Get("Authorization"), qt.Equals, "Basic "+auth) 59 60 c.Assert(r.Body, qt.IsNotNil) 61 defer r.Body.Close() 62 63 body, err := io.ReadAll(r.Body) 64 c.Assert(err, qt.IsNil) 65 c.Check(body, qt.JSONEquals, reqBody) 66 67 w.Header().Set("Content-Type", httpclient.MIMETypeJSON) 68 w.WriteHeader(http.StatusBadRequest) 69 fmt.Fprintln(w, errResp) 70 }) 71 72 srv := httptest.NewServer(h) 73 c.Cleanup(srv.Close) 74 75 exec, err := connector.CreateExecution(nil, cfg(basicAuthType), taskPost) 76 c.Assert(err, qt.IsNil) 77 78 pbIn, err := base.ConvertToStructpb(TaskInput{ 79 EndpointURL: srv.URL + path, 80 Body: reqBody, 81 }) 82 c.Assert(err, qt.IsNil) 83 84 got, err := exec.Execution.Execute([]*structpb.Struct{pbIn}) 85 c.Check(err, qt.IsNil) 86 87 resp := got[0].AsMap() 88 c.Check(resp["status_code"], qt.Equals, float64(http.StatusBadRequest)) 89 c.Check(resp["body"], qt.ContentEquals, map[string]any{"message": "Bad request"}) 90 }) 91 92 c.Run("ok - PUT + query auth", func(c *qt.C) { 93 h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 94 c.Check(r.Method, qt.Equals, http.MethodPut) 95 c.Check(r.URL.Path, qt.Matches, path) 96 97 c.Check(r.FormValue(authKey), qt.Equals, authValue) 98 w.Header().Set("Content-Type", httpclient.MIMETypeJSON) 99 fmt.Fprintln(w, okResp) 100 }) 101 102 srv := httptest.NewServer(h) 103 c.Cleanup(srv.Close) 104 105 exec, err := connector.CreateExecution(nil, cfg(apiKeyType), taskPut) 106 c.Assert(err, qt.IsNil) 107 108 pbIn, err := base.ConvertToStructpb(TaskInput{ 109 EndpointURL: srv.URL + path, 110 Body: reqBody, 111 }) 112 113 c.Assert(err, qt.IsNil) 114 115 got, err := exec.Execution.Execute([]*structpb.Struct{pbIn}) 116 c.Check(err, qt.IsNil) 117 118 resp := got[0].AsMap() 119 c.Check(resp["status_code"], qt.Equals, float64(http.StatusOK)) 120 c.Check(resp["body"], qt.ContentEquals, map[string]any{"title": "Be the wheel"}) 121 }) 122 123 c.Run("ok - GET + bearer auth", func(c *qt.C) { 124 h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 125 c.Check(r.Method, qt.Equals, http.MethodGet) 126 c.Check(r.URL.Path, qt.Matches, path) 127 128 c.Check(r.Header.Get("Authorization"), qt.Equals, "Bearer "+token) 129 130 w.Header().Set("Content-Type", httpclient.MIMETypeJSON) 131 fmt.Fprintln(w, okResp) 132 }) 133 134 srv := httptest.NewServer(h) 135 c.Cleanup(srv.Close) 136 137 exec, err := connector.CreateExecution(nil, cfg(bearerTokenType), taskGet) 138 c.Assert(err, qt.IsNil) 139 140 pbIn, err := base.ConvertToStructpb(TaskInput{ 141 EndpointURL: srv.URL + path, 142 Body: reqBody, 143 }) 144 c.Assert(err, qt.IsNil) 145 146 got, err := exec.Execution.Execute([]*structpb.Struct{pbIn}) 147 c.Check(err, qt.IsNil) 148 149 resp := got[0].AsMap() 150 c.Check(resp["status_code"], qt.Equals, float64(http.StatusOK)) 151 c.Check(resp["body"], qt.ContentEquals, map[string]any{"title": "Be the wheel"}) 152 }) 153 } 154 155 func TestConnector_Test(t *testing.T) { 156 c := qt.New(t) 157 158 logger := zap.NewNop() 159 connector := Init(logger, nil) 160 161 c.Run("ok - connected (even with non-2xx status", func(c *qt.C) { 162 h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 163 c.Check(r.Method, qt.Equals, http.MethodGet) 164 165 w.WriteHeader(http.StatusNotFound) 166 w.Header().Set("Content-Type", "plain/text") 167 fmt.Fprintln(w, "Not Found") 168 }) 169 170 srv := httptest.NewServer(h) 171 c.Cleanup(srv.Close) 172 173 err := connector.Test(nil, cfg(noAuthType)) 174 c.Check(err, qt.IsNil) 175 }) 176 } 177 178 const ( 179 username = "foo" 180 password = "bar" 181 token = "123" 182 authKey = "api-key" 183 authValue = "321" 184 ) 185 186 var testAuth = map[authType]map[string]any{ 187 noAuthType: map[string]any{}, 188 basicAuthType: map[string]any{ 189 "username": username, 190 "password": password, 191 }, 192 bearerTokenType: map[string]any{"token": token}, 193 apiKeyType: map[string]any{ 194 "auth_location": string(query), 195 "key": authKey, 196 "value": authValue, 197 }, 198 } 199 200 func cfg(atype authType) *structpb.Struct { 201 auth := testAuth[atype] 202 auth["auth_type"] = string(atype) 203 config, _ := structpb.NewStruct(map[string]any{ 204 "authentication": auth, 205 }) 206 207 return config 208 }