mirror of https://github.com/matrix-org/go-neb.git
Richard Lewis
8 years ago
3 changed files with 262 additions and 0 deletions
-
1src/github.com/matrix-org/go-neb/goneb.go
-
150src/github.com/matrix-org/go-neb/services/wikipedia/wikipedia.go
-
111src/github.com/matrix-org/go-neb/services/wikipedia/wikipedia_test.go
@ -0,0 +1,150 @@ |
|||||
|
// Package wikipedia implements a Service which adds !commands for Wikipedia search.
|
||||
|
package wikipedia |
||||
|
|
||||
|
import ( |
||||
|
"encoding/json" |
||||
|
"fmt" |
||||
|
"io/ioutil" |
||||
|
"net/http" |
||||
|
"net/url" |
||||
|
"strings" |
||||
|
|
||||
|
log "github.com/Sirupsen/logrus" |
||||
|
"github.com/matrix-org/go-neb/types" |
||||
|
"github.com/matrix-org/gomatrix" |
||||
|
) |
||||
|
|
||||
|
// ServiceType of the Wikipedia service
|
||||
|
const ServiceType = "wikipedia" |
||||
|
|
||||
|
var httpClient = &http.Client{} |
||||
|
|
||||
|
type wikipediaSearchResults struct { |
||||
|
Query struct { |
||||
|
Pages map[string]wikipediaPage `json:"pages"` |
||||
|
} `json:"query"` |
||||
|
} |
||||
|
|
||||
|
type wikipediaPage struct { |
||||
|
PageID int64 `json:"pageid"` |
||||
|
NS int `json:"ns"` |
||||
|
Title string `json:"title"` |
||||
|
Touched string `json:"touched"` |
||||
|
LastRevID int64 `json:"lastrevid"` |
||||
|
Extract string `json:"extract"` |
||||
|
} |
||||
|
|
||||
|
// Service contains the Config fields for the Wikipedia service.
|
||||
|
type Service struct { |
||||
|
types.DefaultService |
||||
|
} |
||||
|
|
||||
|
// Commands supported:
|
||||
|
// !wikipedia some_search_query_without_quotes
|
||||
|
// Responds with a suitable article extract and link to the referenced page into the same room as the command.
|
||||
|
func (s *Service) Commands(client *gomatrix.Client) []types.Command { |
||||
|
return []types.Command{ |
||||
|
types.Command{ |
||||
|
Path: []string{"wikipedia"}, |
||||
|
Command: func(roomID, userID string, args []string) (interface{}, error) { |
||||
|
return s.cmdWikipediaSearch(client, roomID, userID, args) |
||||
|
}, |
||||
|
}, |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// usageMessage returns a matrix TextMessage representation of the service usage
|
||||
|
func usageMessage() *gomatrix.TextMessage { |
||||
|
return &gomatrix.TextMessage{"m.notice", |
||||
|
`Usage: !wikipedia search_text`} |
||||
|
} |
||||
|
|
||||
|
func (s *Service) cmdWikipediaSearch(client *gomatrix.Client, roomID, userID string, args []string) (interface{}, error) { |
||||
|
|
||||
|
if len(args) < 1 { |
||||
|
return usageMessage(), nil |
||||
|
} |
||||
|
|
||||
|
// Get the query text to search for.
|
||||
|
querySentence := strings.Join(args, " ") |
||||
|
|
||||
|
searchResultPage, err := s.text2Wikipedia(querySentence) |
||||
|
|
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
|
||||
|
if searchResultPage.Extract == "" { |
||||
|
return gomatrix.TextMessage{ |
||||
|
MsgType: "m.notice", |
||||
|
Body: "No results found!", |
||||
|
}, nil |
||||
|
} |
||||
|
|
||||
|
return gomatrix.TextMessage{ |
||||
|
MsgType: "m.notice", |
||||
|
Body: searchResultPage.Extract, |
||||
|
}, nil |
||||
|
} |
||||
|
|
||||
|
// text2Wikipedia returns a Wikipedia article summary
|
||||
|
func (s *Service) text2Wikipedia(query string) (*wikipediaPage, error) { |
||||
|
log.Info("Searching Wikipedia for: ", query) |
||||
|
|
||||
|
u, err := url.Parse("https://en.wikipedia.org/w/api.php") |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
|
||||
|
// Example query - https://en.wikipedia.org/w/api.php?action=query&prop=extracts&format=json&exintro=&titles=RMS+Titanic
|
||||
|
q := u.Query() |
||||
|
q.Set("action", "query") // Action - query for articles
|
||||
|
q.Set("prop", "extracts") // Return article extracts
|
||||
|
q.Set("format", "json") |
||||
|
// q.Set("exintro", "")
|
||||
|
q.Set("titles", query) // Text to search for
|
||||
|
|
||||
|
u.RawQuery = q.Encode() |
||||
|
// log.Info("Request URL: ", u)
|
||||
|
|
||||
|
res, err := httpClient.Get(u.String()) |
||||
|
if res != nil { |
||||
|
defer res.Body.Close() |
||||
|
} |
||||
|
if err != nil { |
||||
|
return nil, err |
||||
|
} |
||||
|
if res.StatusCode < 200 || res.StatusCode >= 300 { |
||||
|
return nil, fmt.Errorf("Request error: %d, %s", res.StatusCode, response2String(res)) |
||||
|
} |
||||
|
var searchResults wikipediaSearchResults |
||||
|
|
||||
|
// log.Info(response2String(res))
|
||||
|
if err := json.NewDecoder(res.Body).Decode(&searchResults); err != nil { |
||||
|
return nil, fmt.Errorf("ERROR - %s", err.Error()) |
||||
|
} else if len(searchResults.Pages) < 1 { |
||||
|
return nil, fmt.Errorf("No articles found") |
||||
|
} |
||||
|
|
||||
|
// Return only the first search result
|
||||
|
return &searchResults.Pages[0], nil |
||||
|
} |
||||
|
|
||||
|
// response2String returns a string representation of an HTTP response body
|
||||
|
func response2String(res *http.Response) string { |
||||
|
bs, err := ioutil.ReadAll(res.Body) |
||||
|
if err != nil { |
||||
|
return "Failed to decode response body" |
||||
|
} |
||||
|
str := string(bs) |
||||
|
return str |
||||
|
} |
||||
|
|
||||
|
// Initialise the service
|
||||
|
func init() { |
||||
|
types.RegisterService(func(serviceID, serviceUserID, webhookEndpointURL string) types.Service { |
||||
|
return &Service{ |
||||
|
DefaultService: types.NewDefaultService(serviceID, serviceUserID, ServiceType), |
||||
|
} |
||||
|
}) |
||||
|
} |
@ -0,0 +1,111 @@ |
|||||
|
package wikipedia |
||||
|
|
||||
|
import ( |
||||
|
"bytes" |
||||
|
"encoding/json" |
||||
|
"fmt" |
||||
|
"io/ioutil" |
||||
|
"net/http" |
||||
|
"strings" |
||||
|
"testing" |
||||
|
|
||||
|
"github.com/matrix-org/go-neb/database" |
||||
|
"github.com/matrix-org/go-neb/testutils" |
||||
|
"github.com/matrix-org/go-neb/types" |
||||
|
"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" |
||||
|
wikipediaImageURL := "https://www.wikipediaapis.com/customsearch/v1" |
||||
|
|
||||
|
// Mock the response from Wikipedia
|
||||
|
wikipediaTrans := testutils.NewRoundTripper(func(req *http.Request) (*http.Response, error) { |
||||
|
wikipediaURL := "https://www.wikipediaapis.com/customsearch/v1" |
||||
|
query := req.URL.Query() |
||||
|
|
||||
|
// Check the base API URL
|
||||
|
if !strings.HasPrefix(req.URL.String(), wikipediaURL) { |
||||
|
t.Fatalf("Bad URL: got %s want prefix %s", req.URL.String(), wikipediaURL) |
||||
|
} |
||||
|
// Check the request method
|
||||
|
if req.Method != "GET" { |
||||
|
t.Fatalf("Bad method: got %s want GET", req.Method) |
||||
|
} |
||||
|
// Check the API key
|
||||
|
if query.Get("key") != apiKey { |
||||
|
t.Fatalf("Bad apiKey: got %s want %s", query.Get("key"), apiKey) |
||||
|
} |
||||
|
// Check the search query
|
||||
|
var searchString = query.Get("q") |
||||
|
var searchStringLength = len(searchString) |
||||
|
if searchStringLength > 0 && !strings.HasPrefix(searchString, "image") { |
||||
|
t.Fatalf("Bad search string: got \"%s\" (%d characters) ", searchString, searchStringLength) |
||||
|
} |
||||
|
|
||||
|
resImage := wikipediaImage{ |
||||
|
Width: 64, |
||||
|
Height: 64, |
||||
|
} |
||||
|
|
||||
|
res := wikipediaSearchResult{ |
||||
|
Title: "A Cat", |
||||
|
Link: "http://cat.com/cat.jpg", |
||||
|
Mime: "image/jpeg", |
||||
|
Image: resImage, |
||||
|
} |
||||
|
|
||||
|
b, err := json.Marshal(res) |
||||
|
if err != nil { |
||||
|
t.Fatalf("Failed to marshal Wikipedia response - %s", err) |
||||
|
} |
||||
|
return &http.Response{ |
||||
|
StatusCode: 200, |
||||
|
Body: ioutil.NopCloser(bytes.NewBuffer(b)), |
||||
|
}, nil |
||||
|
}) |
||||
|
// clobber the Wikipedia service http client instance
|
||||
|
httpClient = &http.Client{Transport: wikipediaTrans} |
||||
|
|
||||
|
// Create the Wikipedia service
|
||||
|
srv, err := types.CreateService("id", ServiceType, "@wikipediabot:hyrule", []byte( |
||||
|
`{"api_key":"`+apiKey+`"}`, |
||||
|
)) |
||||
|
if err != nil { |
||||
|
t.Fatal("Failed to create Wikipedia service: ", err) |
||||
|
} |
||||
|
wikipedia := srv.(*Service) |
||||
|
|
||||
|
// Mock the response from Matrix
|
||||
|
matrixTrans := struct{ testutils.MockTransport }{} |
||||
|
matrixTrans.RT = func(req *http.Request) (*http.Response, error) { |
||||
|
if req.URL.String() == wikipediaImageURL { // getting the Wikipedia 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", "@wikipediabot:hyrule", "its_a_secret") |
||||
|
matrixCli.Client = &http.Client{Transport: matrixTrans} |
||||
|
|
||||
|
// Execute the matrix !command
|
||||
|
cmds := wikipedia.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())
|
||||
|
// }
|
||||
|
} |
Write
Preview
Loading…
Cancel
Save
Reference in new issue