weaviate

1
//                           _       _
2
// __      _____  __ ___   ___  __ _| |_ ___
3
// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
4
//  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
5
//   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
6
//
7
//  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
8
//
9
//  CONTACT: hello@weaviate.io
10
//
11

12
package clients
13

14
import (
15
	"context"
16
	"net/http"
17
	"net/http/httptest"
18
	"testing"
19
	"time"
20

21
	"github.com/sirupsen/logrus"
22
	"github.com/sirupsen/logrus/hooks/test"
23
	"github.com/stretchr/testify/assert"
24
	"github.com/stretchr/testify/require"
25
)
26

27
func TestWaitForStartup(t *testing.T) {
28
	t.Run("when the server is immediately ready", func(t *testing.T) {
29
		server := httptest.NewServer(&testReadyHandler{t: t})
30
		defer server.Close()
31
		c := New(server.URL, 0, nullLogger())
32
		err := c.WaitForStartup(context.Background(), 50*time.Millisecond)
33

34
		assert.Nil(t, err)
35
	})
36

37
	t.Run("when the server is down", func(t *testing.T) {
38
		c := New("http://nothing-running-at-this-url", 0, nullLogger())
39
		ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
40
		defer cancel()
41
		err := c.WaitForStartup(ctx, 150*time.Millisecond)
42

43
		require.NotNil(t, err, nullLogger())
44
		assert.Contains(t, err.Error(), "expired before remote was ready")
45
	})
46

47
	t.Run("when the server is alive, but not ready", func(t *testing.T) {
48
		server := httptest.NewServer(&testReadyHandler{
49
			t:         t,
50
			readyTime: time.Now().Add(1 * time.Minute),
51
		})
52
		c := New(server.URL, 0, nullLogger())
53
		defer server.Close()
54
		ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
55
		defer cancel()
56
		err := c.WaitForStartup(ctx, 50*time.Millisecond)
57

58
		require.NotNil(t, err)
59
		assert.Contains(t, err.Error(), "expired before remote was ready")
60
	})
61

62
	t.Run("when the server is initially not ready, but then becomes ready",
63
		func(t *testing.T) {
64
			server := httptest.NewServer(&testReadyHandler{
65
				t:         t,
66
				readyTime: time.Now().Add(100 * time.Millisecond),
67
			})
68
			c := New(server.URL, 0, nullLogger())
69
			defer server.Close()
70
			ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
71
			defer cancel()
72
			err := c.WaitForStartup(ctx, 50*time.Millisecond)
73

74
			require.Nil(t, err)
75
		})
76
}
77

78
type testReadyHandler struct {
79
	t *testing.T
80
	// the test handler will report as not ready before the time has passed
81
	readyTime time.Time
82
}
83

84
func (f *testReadyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
85
	assert.Equal(f.t, "/.well-known/ready", r.URL.String())
86
	assert.Equal(f.t, http.MethodGet, r.Method)
87

88
	if time.Since(f.readyTime) < 0 {
89
		w.WriteHeader(http.StatusServiceUnavailable)
90
	}
91

92
	w.WriteHeader(http.StatusNoContent)
93
}
94

95
func nullLogger() logrus.FieldLogger {
96
	l, _ := test.NewNullLogger()
97
	return l
98
}
99