prysm-pulse/validator/rpc/handlers_health_test.go

package rpc

import (
	"bytes"
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/golang/mock/gomock"
	"github.com/golang/protobuf/ptypes/empty"
	"github.com/prysmaticlabs/prysm/v4/io/logs/mock"
	eth "github.com/prysmaticlabs/prysm/v4/proto/prysm/v1alpha1"
	pb "github.com/prysmaticlabs/prysm/v4/proto/prysm/v1alpha1"
	"github.com/prysmaticlabs/prysm/v4/testing/require"
	validatormock "github.com/prysmaticlabs/prysm/v4/testing/validator-mock"
	"google.golang.org/grpc"
)

type MockBeaconNodeHealthClient struct {
	grpc.ClientStream
	logs []*pb.LogsResponse
	err  error
}

func (m *MockBeaconNodeHealthClient) StreamBeaconLogs(ctx context.Context, in *empty.Empty, opts ...grpc.CallOption) (eth.Health_StreamBeaconLogsClient, error) {
	return m, m.err
}

func (m *MockBeaconNodeHealthClient) Recv() (*eth.LogsResponse, error) {
	if len(m.logs) == 0 {
		return nil, io.EOF
	}
	log := m.logs[0]
	m.logs = m.logs[1:]
	return log, nil
}

func (m *MockBeaconNodeHealthClient) SendMsg(_ interface{}) error {
	return m.err
}

func (m *MockBeaconNodeHealthClient) Context() context.Context {
	return context.Background()
}

type flushableResponseRecorder struct {
	*httptest.ResponseRecorder
	flushed bool
}

func (f *flushableResponseRecorder) Flush() {
	f.flushed = true
}

func TestStreamBeaconLogs(t *testing.T) {
	logs := []*pb.LogsResponse{
		{
			Logs: []string{"log1", "log2"},
		},
		{
			Logs: []string{"log3", "log4"},
		},
	}

	mockClient := &MockBeaconNodeHealthClient{
		logs: logs,
		err:  nil,
	}

	// Setting up the mock in the server struct
	s := Server{
		ctx:                    context.Background(),
		beaconNodeHealthClient: mockClient,
	}

	// Create a mock ResponseWriter and Request
	w := &flushableResponseRecorder{
		ResponseRecorder: httptest.NewRecorder(),
	}
	r := httptest.NewRequest("GET", "/v2/validator/health/logs/beacon/stream", nil)

	// Call the function
	s.StreamBeaconLogs(w, r)

	// Assert the results
	resp := w.Result()
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Expected status OK but got %v", resp.StatusCode)
	}
	ct, ok := resp.Header["Content-Type"]
	require.Equal(t, ok, true)
	require.Equal(t, ct[0], "text/event-stream")
	cn, ok := resp.Header["Connection"]
	require.Equal(t, ok, true)
	require.Equal(t, cn[0], "keep-alive")
	body, err := io.ReadAll(resp.Body)
	require.NoError(t, err)
	require.NotNil(t, body)
	require.StringContains(t, `{"logs":["log1","log2"]}`, string(body))
	require.StringContains(t, `{"logs":["log3","log4"]}`, string(body))
	if !w.flushed {
		t.Fatal("Flush was not called")
	}
}

func TestStreamValidatorLogs(t *testing.T) {
	ctx := context.Background()
	mockLogs := [][]byte{
		[]byte("[2023-10-31 10:00:00] INFO: Starting server..."),
		[]byte("[2023-10-31 10:01:23] DEBUG: Database connection established."),
		[]byte("[2023-10-31 10:05:45] WARN: High memory usage detected."),
		[]byte("[2023-10-31 10:10:12] INFO: New user registered: user123."),
		[]byte("[2023-10-31 10:15:30] ERROR: Failed to send email."),
	}
	logStreamer := mock.NewMockStreamer(mockLogs)
	// Setting up the mock in the server struct
	s := Server{
		ctx:                  ctx,
		logsStreamer:         logStreamer,
		streamLogsBufferSize: 100,
	}

	w := &flushableResponseRecorder{
		ResponseRecorder: httptest.NewRecorder(),
	}
	r := httptest.NewRequest("GET", "/v2/validator/health/logs/validator/stream", nil)
	go func() {
		s.StreamValidatorLogs(w, r)
	}()
	// wait for initiation of StreamValidatorLogs
	time.Sleep(100 * time.Millisecond)
	logStreamer.LogsFeed().Send([]byte("Some mock event data"))
	// wait for feed
	time.Sleep(100 * time.Millisecond)
	s.ctx.Done()
	// Assert the results
	resp := w.Result()
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Expected status OK but got %v", resp.StatusCode)
	}
	ct, ok := resp.Header["Content-Type"]
	require.Equal(t, ok, true)
	require.Equal(t, ct[0], "text/event-stream")
	cn, ok := resp.Header["Connection"]
	require.Equal(t, ok, true)
	require.Equal(t, cn[0], "keep-alive")
	// Check if data was written
	body, err := io.ReadAll(resp.Body)
	require.NoError(t, err)
	require.NotNil(t, body)

	require.StringContains(t, `{"logs":["[2023-10-31 10:00:00] INFO: Starting server...","[2023-10-31 10:01:23] DEBUG: Database connection established.",`+
		`"[2023-10-31 10:05:45] WARN: High memory usage detected.","[2023-10-31 10:10:12] INFO: New user registered: user123.","[2023-10-31 10:15:30] ERROR: Failed to send email."]}`, string(body))
	require.StringContains(t, `{"logs":["Some mock event data"]}`, string(body))

	// Check if Flush was called
	if !w.flushed {
		t.Fatal("Flush was not called")
	}

}

func TestServer_GetVersion(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	ctx := context.Background()
	mockNodeClient := validatormock.NewMockNodeClient(ctrl)
	s := Server{
		ctx:              ctx,
		beaconNodeClient: mockNodeClient,
	}
	mockNodeClient.EXPECT().GetVersion(gomock.Any(), gomock.Any()).Return(&eth.Version{
		Version:  "4.10.1",
		Metadata: "beacon node",
	}, nil)
	r := httptest.NewRequest("GET", "/v2/validator/health/version", nil)
	w := httptest.NewRecorder()
	w.Body = &bytes.Buffer{}
	s.GetVersion(w, r)
	resp := w.Result()
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Expected status OK but got %v", resp.StatusCode)
	}
	body, err := io.ReadAll(resp.Body)
	require.NoError(t, err)
	require.NotNil(t, body)
	require.StringContains(t, `{"beacon":"4.10.1","validator":"Prysm/Unknown/Local build. Built at: Moments ago"}`, string(body))
}