/
google_screenshot.go
123 lines (103 loc) · 3.2 KB
/
google_screenshot.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
// Copyright 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// Package googlescreenshot includes a handler for an advanced screenshot endpoint at POST google/screenshot.
package googlescreenshot
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"image"
"image/png"
"log"
"net/http"
"github.com/bazelbuild/rules_webtesting/go/cropper"
"github.com/bazelbuild/rules_webtesting/go/metadata/capabilities"
"github.com/bazelbuild/rules_webtesting/go/webdriver"
"github.com/bazelbuild/rules_webtesting/go/wtl/proxy/driverhub"
)
type request struct {
Element map[string]interface{}
Exclude []map[string]interface{}
}
// ProviderFunc provides a handler for an advanced screenshot endpoint at POST google/screenshot.
func ProviderFunc(session *driverhub.WebDriverSession, _ *capabilities.Capabilities, base driverhub.HandlerFunc) (driverhub.HandlerFunc, bool) {
return func(ctx context.Context, rq driverhub.Request) (driverhub.Response, error) {
if rq.Method != http.MethodPost || len(rq.Path) != 2 || rq.Path[0] != "google" || rq.Path[1] != "screenshot" {
return base(ctx, rq)
}
r := request{}
if err := json.Unmarshal(rq.Body, &r); err != nil {
// TODO(DrMarcII): better error message?
return driverhub.ResponseFromError(err)
}
img, err := captureScreenshot(ctx, session.WebDriver, r)
if err != nil {
// TODO(DrMarcII): better error message?
return driverhub.ResponseFromError(err)
}
return createResponse(img)
}, true
}
func captureScreenshot(ctx context.Context, driver webdriver.WebDriver, r request) (image.Image, error) {
el, _ := driver.ElementFromMap(r.Element)
if el != nil {
if err := el.ScrollIntoView(ctx); err != nil {
return nil, err
}
}
// TODO(DrMarcII): stabilization to ensure that we have finished scrolling?
img, err := driver.Screenshot(ctx)
if err != nil {
return nil, err
}
for _, e := range r.Exclude {
exclude, err := driver.ElementFromMap(e)
if err != nil {
log.Print(err)
continue
}
bounds, err := exclude.Bounds(ctx)
if err != nil {
log.Print(err)
continue
}
i, err := cropper.Blackout(img, bounds)
if err != nil {
return nil, err
}
img = i
}
if el != nil {
bounds, err := el.Bounds(ctx)
if err != nil {
return nil, err
}
return cropper.Crop(img, bounds)
}
return img, nil
}
func createResponse(img image.Image) (driverhub.Response, error) {
buffer := &bytes.Buffer{}
b64 := base64.NewEncoder(base64.StdEncoding, buffer)
defer b64.Close()
if err := png.Encode(b64, img); err != nil {
return driverhub.ResponseFromError(err)
}
return driverhub.Response{
Status: http.StatusOK,
Body: []byte(fmt.Sprintf(`{"status": 0, "value": %q}`, buffer.String())),
}, nil
}