github.com/google/osv-scalibr@v0.4.1/clients/clienttest/mock_http.go (about)

     1  // Copyright 2025 Google LLC
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package clienttest provides mock servers for testing.
    16  package clienttest
    17  
    18  import (
    19  	"log"
    20  	"net/http"
    21  	"net/http/httptest"
    22  	"os"
    23  	"strings"
    24  	"sync"
    25  	"testing"
    26  )
    27  
    28  // MockHTTPServer is a simple HTTP Server for mocking basic requests.
    29  type MockHTTPServer struct {
    30  	*httptest.Server
    31  
    32  	mu            sync.Mutex
    33  	response      map[string][]byte // path -> response
    34  	authorization string            // expected Authorization header contents
    35  }
    36  
    37  // NewMockHTTPServer starts and returns a new simple HTTP Server for mocking basic requests.
    38  // The Server will automatically be shut down with Close() in the test Cleanup function.
    39  //
    40  // Use the SetResponse / SetResponseFromFile to set the responses for specific URL paths.
    41  func NewMockHTTPServer(t *testing.T) *MockHTTPServer {
    42  	t.Helper()
    43  	mock := &MockHTTPServer{response: make(map[string][]byte)}
    44  	mock.Server = httptest.NewServer(mock)
    45  	t.Cleanup(func() { mock.Close() })
    46  
    47  	return mock
    48  }
    49  
    50  // SetResponse sets the Server's response for the URL path to be response bytes.
    51  func (m *MockHTTPServer) SetResponse(t *testing.T, path string, response []byte) {
    52  	t.Helper()
    53  	m.mu.Lock()
    54  	defer m.mu.Unlock()
    55  	path = strings.TrimPrefix(path, "/")
    56  	m.response[path] = response
    57  }
    58  
    59  // SetResponseFromFile sets the Server's response for the URL path to be the contents of the file at filename.
    60  func (m *MockHTTPServer) SetResponseFromFile(t *testing.T, path string, filename string) {
    61  	t.Helper()
    62  	b, err := os.ReadFile(filename)
    63  	if err != nil {
    64  		t.Fatalf("failed to read response file: %v", err)
    65  	}
    66  	m.SetResponse(t, path, b)
    67  }
    68  
    69  // SetAuthorization sets the contents of the 'Authorization' header the server expects for all endpoints.
    70  //
    71  // The incoming requests' headers must match the auth string exactly, otherwise the server will response with 401 Unauthorized.
    72  // If authorization is unset or empty, the server will not require authorization.
    73  func (m *MockHTTPServer) SetAuthorization(t *testing.T, auth string) {
    74  	t.Helper()
    75  	m.mu.Lock()
    76  	defer m.mu.Unlock()
    77  	m.authorization = auth
    78  }
    79  
    80  // ServeHTTP is the http.Handler for the underlying httptest.Server.
    81  func (m *MockHTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    82  	m.mu.Lock()
    83  	wantAuth := m.authorization
    84  	resp, ok := m.response[strings.TrimPrefix(r.URL.EscapedPath(), "/")]
    85  	m.mu.Unlock()
    86  
    87  	if wantAuth != "" && r.Header.Get("Authorization") != wantAuth {
    88  		w.WriteHeader(http.StatusUnauthorized)
    89  		resp = []byte("unauthorized")
    90  	} else if !ok {
    91  		w.WriteHeader(http.StatusNotFound)
    92  		resp = []byte("not found")
    93  	}
    94  
    95  	if _, err := w.Write(resp); err != nil {
    96  		log.Fatalf("Write: %v", err)
    97  	}
    98  }