From 2f49617eb6ec85126a0fe7a7e2e3812673324571 Mon Sep 17 00:00:00 2001 From: Michael Telatynski <7t3chguy@gmail.com> Date: Thu, 27 Jun 2019 20:02:16 +0100 Subject: [PATCH] Add safe-search option to google service --- config.sample.yaml | 1 + .../go-neb/services/google/google.go | 8 ++ .../go-neb/services/google/google_test.go | 81 ++++++++++++++++--- 3 files changed, 81 insertions(+), 9 deletions(-) diff --git a/config.sample.yaml b/config.sample.yaml index dc196c9..bbf9afc 100644 --- a/config.sample.yaml +++ b/config.sample.yaml @@ -81,6 +81,7 @@ services: Config: api_key: "AIzaSyA4FD39m9" cx: "AIASDFWSRRtrtr" + safe_search: false - ID: "imgur_service" Type: "imgur" diff --git a/src/github.com/matrix-org/go-neb/services/google/google.go b/src/github.com/matrix-org/go-neb/services/google/google.go index 347011b..07a9dee 100644 --- a/src/github.com/matrix-org/go-neb/services/google/google.go +++ b/src/github.com/matrix-org/go-neb/services/google/google.go @@ -63,6 +63,8 @@ type Service struct { APIKey string `json:"api_key"` // The Google custom search engine ID Cx string `json:"cx"` + // Whether or not to enable Google safe-search + Safesearch bool `json:"safe_search"` } // Commands supported: @@ -157,6 +159,12 @@ func (s *Service) text2imgGoogle(query string) (*googleSearchResult, error) { q.Set("key", s.APIKey) // Set the API key for the request q.Set("cx", s.Cx) // Set the custom search engine ID + if s.Safesearch { + q.Set("safe", "active") + } else { + q.Set("safe", "off") + } + u.RawQuery = q.Encode() // log.Info("Request URL: ", u) diff --git a/src/github.com/matrix-org/go-neb/services/google/google_test.go b/src/github.com/matrix-org/go-neb/services/google/google_test.go index d3a5cf7..a888979 100644 --- a/src/github.com/matrix-org/go-neb/services/google/google_test.go +++ b/src/github.com/matrix-org/go-neb/services/google/google_test.go @@ -15,15 +15,8 @@ import ( "github.com/matrix-org/gomatrix" ) -// TODO: It would be nice to tabularise this test so we can try failing different combinations of responses to make -// sure all cases are handled, rather than just the general case as is here. -func TestCommand(t *testing.T) { - database.SetServiceDB(&database.NopStorage{}) - apiKey := "secret" - googleImageURL := "http://cat.com/cat.jpg" - - // Mock the response from Google - googleTrans := testutils.NewRoundTripper(func(req *http.Request) (*http.Response, error) { +func mockGoogle(t *testing.T, apiKey, googleImageURL string, safeSearch bool) http.RoundTripper { + return testutils.NewRoundTripper(func(req *http.Request) (*http.Response, error) { googleURL := "https://www.googleapis.com/customsearch/v1" query := req.URL.Query() @@ -39,6 +32,15 @@ func TestCommand(t *testing.T) { if query.Get("key") != apiKey { t.Fatalf("Bad apiKey: got %s want %s", query.Get("key"), apiKey) } + // Check safe-search + safe := "off" + if safeSearch { + safe = "active" + } + if query.Get("safe") != safe { + t.Fatalf("Bad safe: got %s want %s", query.Get("safe"), safe) + } + // Check the search query var searchString = query.Get("q") var searchStringLength = len(searchString) @@ -73,6 +75,17 @@ func TestCommand(t *testing.T) { Body: ioutil.NopCloser(bytes.NewBuffer(b)), }, nil }) +} + +// TODO: It would be nice to tabularise this test so we can try failing different combinations of responses to make +// sure all cases are handled, rather than just the general case as is here. +func TestCommand(t *testing.T) { + database.SetServiceDB(&database.NopStorage{}) + apiKey := "secret" + googleImageURL := "http://cat.com/cat.jpg" + + // Mock the response from Google + googleTrans := mockGoogle(t, apiKey, googleImageURL, false) // clobber the Google service http client instance httpClient = &http.Client{Transport: googleTrans} @@ -115,3 +128,53 @@ func TestCommand(t *testing.T) { t.Fatalf("Failed to process command: %s", err.Error()) } } + +func TestSafesearch(t *testing.T) { + database.SetServiceDB(&database.NopStorage{}) + apiKey := "secret" + googleImageURL := "http://cat.com/cat.jpg" + + // Mock the response from Google + googleTrans := mockGoogle(t, apiKey, googleImageURL, true) + // clobber the Google service http client instance + httpClient = &http.Client{Transport: googleTrans} + + // Create the Google service + srv, err := types.CreateService("id", ServiceType, "@googlebot:hyrule", []byte( + `{"api_key":"`+apiKey+`", "safe_search": true}`, + )) + if err != nil { + t.Fatal("Failed to create Google service: ", err) + } + google := srv.(*Service) + + // Mock the response from Matrix + matrixTrans := struct{ testutils.MockTransport }{} + matrixTrans.RT = func(req *http.Request) (*http.Response, error) { + if req.URL.String() == googleImageURL { // getting the Google image + return &http.Response{ + StatusCode: 200, + Body: ioutil.NopCloser(bytes.NewBufferString("some image data")), + }, nil + } else if strings.Contains(req.URL.String(), "_matrix/media/r0/upload") { // uploading the image to matrix + return &http.Response{ + StatusCode: 200, + Body: ioutil.NopCloser(bytes.NewBufferString(`{"content_uri":"mxc://foo/bar"}`)), + }, nil + } + return nil, fmt.Errorf("Unknown URL: %s", req.URL.String()) + } + matrixCli, _ := gomatrix.NewClient("https://hyrule", "@googlebot:hyrule", "its_a_secret") + matrixCli.Client = &http.Client{Transport: matrixTrans} + + // Execute the matrix !command + cmds := google.Commands(matrixCli) + if len(cmds) != 3 { + t.Fatalf("Unexpected number of commands: %d", len(cmds)) + } + cmd := cmds[0] + _, err = cmd.Command("!someroom:hyrule", "@navi:hyrule", []string{"image", "Czechoslovakian bananna"}) + if err != nil { + t.Fatalf("Failed to process command: %s", err.Error()) + } +}