github.com/filecoin-project/bacalhau@v0.3.23-0.20230228154132-45c989550ace/pkg/test/requester/publicapi/websocket_test.go (about)

     1  //go:build unit || !integration
     2  
     3  package publicapi
     4  
     5  import (
     6  	"context"
     7  	"fmt"
     8  	"testing"
     9  
    10  	"github.com/filecoin-project/bacalhau/pkg/logger"
    11  	"github.com/filecoin-project/bacalhau/pkg/model"
    12  	"github.com/filecoin-project/bacalhau/pkg/node"
    13  	"github.com/filecoin-project/bacalhau/pkg/requester/publicapi"
    14  	testutils "github.com/filecoin-project/bacalhau/pkg/test/utils"
    15  	"github.com/gorilla/websocket"
    16  	"github.com/stretchr/testify/require"
    17  	"github.com/stretchr/testify/suite"
    18  )
    19  
    20  // Define the suite, and absorb the built-in basic suite
    21  // functionality from testify - including a T() method which
    22  // returns the current testing context
    23  type WebsocketSuite struct {
    24  	suite.Suite
    25  	node   *node.Node
    26  	client *publicapi.RequesterAPIClient
    27  }
    28  
    29  // In order for 'go test' to run this suite, we need to create
    30  // a normal test function and pass our suite to suite.Run
    31  func TestWebsocketSuite(t *testing.T) {
    32  	suite.Run(t, new(WebsocketSuite))
    33  }
    34  
    35  // Before each test
    36  func (s *WebsocketSuite) SetupTest() {
    37  	logger.ConfigureTestLogging(s.T())
    38  	n, client := setupNodeForTest(s.T())
    39  	s.node = n
    40  	s.client = client
    41  }
    42  
    43  // After each test
    44  func (s *WebsocketSuite) TearDownTest() {
    45  	s.node.CleanupManager.Cleanup(context.Background())
    46  }
    47  
    48  func (s *WebsocketSuite) TestWebsocketEverything() {
    49  	ctx := context.Background()
    50  	// string.Replace http with ws in c.BaseURI
    51  	url := "ws" + s.client.BaseURI[4:] + "/requester/websocket/events"
    52  
    53  	conn, _, err := websocket.DefaultDialer.Dial(url, nil)
    54  	require.NoError(s.T(), err)
    55  
    56  	eventChan := make(chan model.JobEvent)
    57  	go func() {
    58  		for {
    59  			var event model.JobEvent
    60  			err = conn.ReadJSON(&event)
    61  			require.NoError(s.T(), err)
    62  			eventChan <- event
    63  		}
    64  	}()
    65  
    66  	genericJob := testutils.MakeGenericJob()
    67  	_, err = s.client.Submit(ctx, genericJob)
    68  	require.NoError(s.T(), err)
    69  
    70  	event := <-eventChan
    71  	require.Equal(s.T(), "Created", event.EventName.String())
    72  
    73  }
    74  
    75  func (s *WebsocketSuite) TestWebsocketSingleJob() {
    76  	s.T().Skip("TODO: test is flaky as by the time we connect to the websocket, " +
    77  		"the job has already progressed and first event is not guaranteed to be 'Created'")
    78  	ctx := context.Background()
    79  
    80  	genericJob := testutils.MakeGenericJob()
    81  	j, err := s.client.Submit(ctx, genericJob)
    82  	require.NoError(s.T(), err)
    83  
    84  	url := "ws" + s.client.BaseURI[4:] + fmt.Sprintf("/websocket/events?job_id=%s", j.Metadata.ID)
    85  	conn, _, err := websocket.DefaultDialer.Dial(url, nil)
    86  	require.NoError(s.T(), err)
    87  
    88  	var event model.JobEvent
    89  	err = conn.ReadJSON(&event)
    90  	require.NoError(s.T(), err)
    91  	require.Equal(s.T(), "Created", event.EventName.String())
    92  }