Repository: cli/cli
File: internal/codespaces/connection/tunnels_api_server_mock.go (Go)

Code before:
	"log"
	"net/http"
	"net/http/httptest"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/microsoft/dev-tunnels/go/tunnels"
	tunnelssh "github.com/microsoft/dev-tunnels/go/tunnels/ssh"
	"github.com/microsoft/dev-tunnels/go/tunnels/ssh/messages"
	"golang.org/x/crypto/ssh"
)

func NewMockHttpClient() (*http.Client, error) {
	accessToken := "tunnel access-token"
	relayServer, err := newMockrelayServer(withAccessToken(accessToken))
	if err != nil {
		return nil, fmt.Errorf("NewrelayServer returned an error: %w", err)
	}

	hostURL := strings.Replace(relayServer.URL(), "http://", "ws://", 1)
	mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var response []byte
		if r.URL.Path == "/api/v1/tunnels/tunnel-id" {
			tunnel := &tunnels.Tunnel{
				AccessTokens: map[tunnels.TunnelAccessScope]string{
					tunnels.TunnelAccessScopeConnect: accessToken,
				},
				Endpoints: []tunnels.TunnelEndpoint{
					{
						HostID: "host1",
						TunnelRelayTunnelEndpoint: tunnels.TunnelRelayTunnelEndpoint{
							ClientRelayURI: hostURL,
						},
					},
				},
			}

			response, err = json.Marshal(*tunnel)
			if err != nil {
				log.Fatalf("json.Marshal returned an error: %v", err)
			}
		} else if strings.HasPrefix(r.URL.Path, "/api/v1/tunnels/tunnel-id/ports") {
			// Use regex to check if the path ends with a number
			match, err := regexp.MatchString(`\/\d+$`, r.URL.Path)
			if err != nil {
				log.Fatalf("regexp.MatchString returned an error: %v", err)
			}

Review comment (suggestion):
The majority of changes in this file relate to this. Previous behaviour was maintained when this is not provided to avoid a lot of test churn. If you provide these ports, they will be used when working with the http server tunnel port endpoints.

Code after:
	"log"
	"net/http"
	"net/http/httptest"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/microsoft/dev-tunnels/go/tunnels"
	tunnelssh "github.com/microsoft/dev-tunnels/go/tunnels/ssh"
	"github.com/microsoft/dev-tunnels/go/tunnels/ssh/messages"
	"golang.org/x/crypto/ssh"
)

type mockClientOpts struct {
	ports map[int]tunnels.TunnelPort // Port number to protocol
}

type mockClientOpt func(*mockClientOpts)

// WithSpecificPorts allows you to specify a map of ports to TunnelPorts that will be returned by the mock HTTP client.
// Note that this does not take a copy of the map, so you should not modify the map after passing it to this function.
func WithSpecificPorts(ports map[int]tunnels.TunnelPort) mockClientOpt {
	return func(opts *mockClientOpts) {
		opts.ports = ports
	}
}

func NewMockHttpClient(opts ...mockClientOpt) (*http.Client, error) {
	mockClientOpts := &mockClientOpts{}
	for _, opt := range opts {
		opt(mockClientOpts)
	}

	specifiedPorts := mockClientOpts.ports

	accessToken := "tunnel access-token"
	relayServer, err := newMockrelayServer(withAccessToken(accessToken))
	if err != nil {
		return nil, fmt.Errorf("NewrelayServer returned an error: %w", err)
	}

	hostURL := strings.Replace(relayServer.URL(), "http://", "ws://", 1)
	mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var response []byte
		if r.URL.Path == "/tunnels/tunnel-id" {
			tunnel := &tunnels.Tunnel{
				AccessTokens: map[tunnels.TunnelAccessScope]string{