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  }