From 70bc4927791002677b6b1425c48bc76021eb4ff1 Mon Sep 17 00:00:00 2001
From: Samuel Hawksby-Robinson <samuel@samyoul.com>
Date: Wed, 12 Oct 2022 13:51:13 +0100
Subject: [PATCH] Added functionality to ensure that Make...URL() funcs wait
 until the port is ready. Adapted tests to check for wait.

---
 server/ports.go       |  4 ++++
 server/server.go      |  9 ++++++---
 server/server_test.go | 35 +++++++++++++++++++++++++++--------
 3 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/server/ports.go b/server/ports.go
index 01e854a02..b0c57bb0b 100644
--- a/server/ports.go
+++ b/server/ports.go
@@ -36,6 +36,10 @@ func (p *portManger) ResetPort() {
 	}
 }
 
+func (p *portManger) GetPort() int {
+	return p.port
+}
+
 func (p *portManger) MustGetPort() int {
 	p.portWait.Lock()
 	defer p.portWait.Unlock()
diff --git a/server/server.go b/server/server.go
index 744ec6a11..7e8a09df2 100644
--- a/server/server.go
+++ b/server/server.go
@@ -33,8 +33,11 @@ func NewServer(cert *tls.Certificate, hostname string, afterPortChanged func(int
 }
 
 func (s *Server) getHost() string {
-	// TODO consider returning an error if s.getPort returns `0`, as this means that the listener is not ready
-	return fmt.Sprintf("%s:%d", s.hostname, s.port)
+	return fmt.Sprintf("%s:%d", s.hostname, s.GetPort())
+}
+
+func (s *Server) mustGetHost() string {
+	return fmt.Sprintf("%s:%d", s.hostname, s.MustGetPort())
 }
 
 func (s *Server) listenAndServe() {
@@ -145,6 +148,6 @@ func (s *Server) AddHandlers(handlers HandlerPatternMap) {
 func (s *Server) MakeBaseURL() *url.URL {
 	return &url.URL{
 		Scheme: "https",
-		Host:   s.getHost(),
+		Host:   s.mustGetHost(),
 	}
 }
diff --git a/server/server_test.go b/server/server_test.go
index 783d5a94b..21e86c055 100644
--- a/server/server_test.go
+++ b/server/server_test.go
@@ -2,10 +2,15 @@ package server
 
 import (
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/suite"
 )
 
+const (
+	waitTime = 50 * time.Millisecond
+)
+
 func TestServerURLSuite(t *testing.T) {
 	suite.Run(t, new(ServerURLSuite))
 }
@@ -16,9 +21,10 @@ type ServerURLSuite struct {
 
 	server       *MediaServer
 	serverNoPort *MediaServer
+	testStart    time.Time
 }
 
-func (s *ServerURLSuite) SetupSuite() {
+func (s *ServerURLSuite) SetupTest() {
 	s.SetupKeyComponents(s.T())
 
 	s.server = &MediaServer{Server: Server{
@@ -29,25 +35,38 @@ func (s *ServerURLSuite) SetupSuite() {
 	s.Require().NoError(err)
 
 	s.serverNoPort = &MediaServer{Server: Server{
-		hostname: defaultIP.String(),
+		hostname:   defaultIP.String(),
+		portManger: newPortManager(nil),
 	}}
+	go func() {
+		time.Sleep(waitTime)
+		s.serverNoPort.port = 0
+		s.serverNoPort.portWait.Unlock()
+	}()
+
+	s.testStart = time.Now()
+}
+
+func (s *ServerURLSuite) testNoPort(expected string, actual string) {
+	s.Require().Equal(expected, actual)
+	s.Require().Greater(time.Now().Sub(s.testStart), waitTime)
 }
 
 func (s *ServerURLSuite) TestServer_MakeBaseURL() {
 	s.Require().Equal("https://127.0.0.1:1337", s.server.MakeBaseURL().String())
-	s.Require().Equal("https://127.0.0.1:0", s.serverNoPort.MakeBaseURL().String())
+	s.testNoPort("https://127.0.0.1:0", s.serverNoPort.MakeBaseURL().String())
 }
 
 func (s *ServerURLSuite) TestServer_MakeImageServerURL() {
 	s.Require().Equal("https://127.0.0.1:1337/messages/", s.server.MakeImageServerURL())
-	s.Require().Equal("https://127.0.0.1:0/messages/", s.serverNoPort.MakeImageServerURL())
+	s.testNoPort("https://127.0.0.1:0/messages/", s.serverNoPort.MakeImageServerURL())
 }
 
 func (s *ServerURLSuite) TestServer_MakeIdenticonURL() {
 	s.Require().Equal(
 		"https://127.0.0.1:1337/messages/identicons?publicKey=0xdaff0d11decade",
 		s.server.MakeIdenticonURL("0xdaff0d11decade"))
-	s.Require().Equal(
+	s.testNoPort(
 		"https://127.0.0.1:0/messages/identicons?publicKey=0xdaff0d11decade",
 		s.serverNoPort.MakeIdenticonURL("0xdaff0d11decade"))
 }
@@ -56,7 +75,7 @@ func (s *ServerURLSuite) TestServer_MakeImageURL() {
 	s.Require().Equal(
 		"https://127.0.0.1:1337/messages/images?messageId=0x10aded70ffee",
 		s.server.MakeImageURL("0x10aded70ffee"))
-	s.Require().Equal(
+	s.testNoPort(
 		"https://127.0.0.1:0/messages/images?messageId=0x10aded70ffee",
 		s.serverNoPort.MakeImageURL("0x10aded70ffee"))
 }
@@ -65,7 +84,7 @@ func (s *ServerURLSuite) TestServer_MakeAudioURL() {
 	s.Require().Equal(
 		"https://127.0.0.1:1337/messages/audio?messageId=0xde1e7ebee71e",
 		s.server.MakeAudioURL("0xde1e7ebee71e"))
-	s.Require().Equal(
+	s.testNoPort(
 		"https://127.0.0.1:0/messages/audio?messageId=0xde1e7ebee71e",
 		s.serverNoPort.MakeAudioURL("0xde1e7ebee71e"))
 }
@@ -74,7 +93,7 @@ func (s *ServerURLSuite) TestServer_MakeStickerURL() {
 	s.Require().Equal(
 		"https://127.0.0.1:1337/ipfs?hash=0xdeadbeef4ac0",
 		s.server.MakeStickerURL("0xdeadbeef4ac0"))
-	s.Require().Equal(
+	s.testNoPort(
 		"https://127.0.0.1:0/ipfs?hash=0xdeadbeef4ac0",
 		s.serverNoPort.MakeStickerURL("0xdeadbeef4ac0"))
 }