diff --git a/tw/twdataloader/helpers.go b/tw/twdataloader/helpers.go index 27f25f8..c12e576 100644 --- a/tw/twdataloader/helpers.go +++ b/tw/twdataloader/helpers.go @@ -5,6 +5,7 @@ import ( "encoding/csv" "io" "net/http" + "net/http/httptest" "time" ) @@ -22,3 +23,41 @@ func uncompressAndReadCsvLines(r io.Reader) ([][]string, error) { defer uncompressedStream.Close() return csv.NewReader(uncompressedStream).ReadAll() } + +type handlers struct { + getServers http.HandlerFunc +} + +func (h *handlers) init() { + noop := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + }) + if h.getServers == nil { + h.getServers = noop + } +} + +func prepareTestServer(h *handlers) *httptest.Server { + if h == nil { + h = &handlers{} + } + h.init() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case EndpointGetServers: + h.getServers(w, r) + return + default: + w.WriteHeader(http.StatusNotFound) + } + })) +} + +func createWriteStringHandler(resp string) http.HandlerFunc { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, err := w.Write([]byte(resp)) + if err != nil { + w.WriteHeader(http.StatusInternalServerError) + } + }) +} diff --git a/tw/twdataloader/version_data_loader_test.go b/tw/twdataloader/version_data_loader_test.go index ef015e2..aa69808 100644 --- a/tw/twdataloader/version_data_loader_test.go +++ b/tw/twdataloader/version_data_loader_test.go @@ -2,28 +2,11 @@ package twdataloader import ( "github.com/stretchr/testify/assert" - "net/http" - "net/http/httptest" "strconv" "strings" "testing" ) -func prepareTestServer(resp string) *httptest.Server { - return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case EndpointGetServers: - _, err := w.Write([]byte(resp)) - if err != nil { - w.WriteHeader(http.StatusInternalServerError) - } - return - default: - w.WriteHeader(http.StatusNotFound) - } - })) -} - func TestLoadServers(t *testing.T) { t.Run("invalid response", func(t *testing.T) { type scenario struct { @@ -55,7 +38,9 @@ func TestLoadServers(t *testing.T) { } for _, scenario := range scenarios { - ts := prepareTestServer(scenario.resp) + ts := prepareTestServer(&handlers{ + getServers: createWriteStringHandler(scenario.resp), + }) dl := NewVersionDataLoader(&VersionDataLoaderConfig{ Host: strings.ReplaceAll(ts.URL, "https://", ""), @@ -72,7 +57,9 @@ func TestLoadServers(t *testing.T) { t.Run("success", func(t *testing.T) { resp := `a:19:{s:5:"pl150";s:25:"https://pl150.plemiona.pl";s:5:"pl151";s:25:"https://pl151.plemiona.pl";s:5:"pl152";s:25:"https://pl152.plemiona.pl";s:5:"pl153";s:25:"https://pl153.plemiona.pl";s:5:"pl154";s:25:"https://pl154.plemiona.pl";s:5:"pl155";s:25:"https://pl155.plemiona.pl";s:5:"pl156";s:25:"https://pl156.plemiona.pl";s:5:"pl157";s:25:"https://pl157.plemiona.pl";s:5:"pl158";s:25:"https://pl158.plemiona.pl";s:5:"pl159";s:25:"https://pl159.plemiona.pl";s:5:"pl160";s:25:"https://pl160.plemiona.pl";s:5:"pl161";s:25:"https://pl161.plemiona.pl";s:5:"pl162";s:25:"https://pl162.plemiona.pl";s:5:"pl163";s:25:"https://pl163.plemiona.pl";s:5:"pl164";s:25:"https://pl164.plemiona.pl";s:4:"plp7";s:24:"https://plp7.plemiona.pl";s:5:"pl165";s:25:"https://pl165.plemiona.pl";s:5:"pl166";s:25:"https://pl166.plemiona.pl";s:5:"pl167";s:25:"https://pl167.plemiona.pl";}` - ts := prepareTestServer(resp) + ts := prepareTestServer(&handlers{ + getServers: createWriteStringHandler(resp), + }) defer ts.Close() expectedLength, err := strconv.Atoi(strings.Split(resp, ":")[1])