From 2cd364224fddadbcf5c58d82de564930643b10f6 Mon Sep 17 00:00:00 2001
From: Kegan Dougal <7190048+kegsay@users.noreply.github.com>
Date: Thu, 30 Nov 2023 09:55:39 +0000
Subject: [PATCH] Always use mitmproxy in tests FOR CLIENTS

---
 internal/deploy/deploy.go | 19 +++++++++++++++++++
 tests/main_test.go        |  4 +++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/internal/deploy/deploy.go b/internal/deploy/deploy.go
index d717392..773ccfc 100644
--- a/internal/deploy/deploy.go
+++ b/internal/deploy/deploy.go
@@ -8,6 +8,7 @@ import (
 	"net/url"
 	"os"
 	"os/exec"
+	"sync"
 	"testing"
 	"time"
 
@@ -24,6 +25,8 @@ type SlidingSyncDeployment struct {
 	slidingSync    testcontainers.Container
 	reverseProxy   testcontainers.Container
 	slidingSyncURL string
+	proxyURLToHS   map[string]string
+	mu             sync.RWMutex
 	tcpdump        *exec.Cmd
 }
 
@@ -36,6 +39,12 @@ func (d *SlidingSyncDeployment) SlidingSyncURL(t *testing.T) string {
 	return d.slidingSyncURL
 }
 
+func (d *SlidingSyncDeployment) ReverseProxyURLForHS(hsName string) string {
+	d.mu.RLock()
+	defer d.mu.RUnlock()
+	return d.proxyURLToHS[hsName]
+}
+
 func (d *SlidingSyncDeployment) Teardown(writeLogs bool) {
 	if d.slidingSync != nil {
 		if writeLogs {
@@ -54,6 +63,12 @@ func (d *SlidingSyncDeployment) Teardown(writeLogs bool) {
 		}
 	}
 	if d.reverseProxy != nil {
+		if writeLogs {
+			err := writeContainerLogs(d.reverseProxy, "container-mitmproxy.log")
+			if err != nil {
+				log.Printf("failed to write sliding sync logs: %s", err)
+			}
+		}
 		if err := d.reverseProxy.Terminate(context.Background()); err != nil {
 			log.Fatalf("failed to stop reverse proxy: %s", err)
 		}
@@ -176,6 +191,10 @@ func RunNewDeployment(t *testing.T, shouldTCPDump bool) *SlidingSyncDeployment {
 		reverseProxy:   mitmproxyContainer,
 		slidingSyncURL: ssURL,
 		tcpdump:        cmd,
+		proxyURLToHS: map[string]string{
+			"hs1": rpHS1URL,
+			"hs2": rpHS2URL,
+		},
 	}
 }
 
diff --git a/tests/main_test.go b/tests/main_test.go
index a228a5b..d08521b 100644
--- a/tests/main_test.go
+++ b/tests/main_test.go
@@ -158,5 +158,7 @@ func (c *TestContext) MustLoginDevice(t *testing.T, existing *client.CSAPI, clie
 
 func (c *TestContext) MustLoginClient(t *testing.T, cli *client.CSAPI, clientType api.ClientType) api.Client {
 	t.Helper()
-	return MustLoginClient(t, clientType, api.FromComplementClient(cli, "complement-crypto-password"), c.Deployment.SlidingSyncURL(t))
+	cfg := api.FromComplementClient(cli, "complement-crypto-password")
+	cfg.BaseURL = c.Deployment.ReverseProxyURLForHS(clientType.HS)
+	return MustLoginClient(t, clientType, cfg, c.Deployment.SlidingSyncURL(t))
 }