[cleanup] Enable Beacon assertions in Tests

tavern/internal/c2/api_claim_tasks_test.go

@@ -3,15 +3,18 @@ package c2_test
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/testing/protocmp"
 	"realm.pub/tavern/internal/c2/c2pb"
 	"realm.pub/tavern/internal/c2/c2test"
 	"realm.pub/tavern/internal/ent"
+	"realm.pub/tavern/internal/ent/beacon"
 )
 
 func TestClaimTasks(t *testing.T) {
@@ -28,13 +31,16 @@ func TestClaimTasks(t *testing.T) {
 	}
 
 	// Test Cases
-	type testCase struct {
+	tests := []struct {
 		name     string
 		req      *c2pb.ClaimTasksRequest
 		wantResp *c2pb.ClaimTasksResponse
 		wantCode codes.Code
-	}
-	tests := []testCase{
+
+		wantBeaconExist            bool
+		wantBeaconLastSeenAtBefore time.Time
+		wantBeaconLastSeenAtAfter  time.Time
+	}{
 		{
 			name: "First_Callback",
 			req: &c2pb.ClaimTasksRequest{
@@ -50,11 +56,15 @@ func TestClaimTasks(t *testing.T) {
 						Platform:   c2pb.Host_PLATFORM_LINUX,
 						PrimaryIp:  "127.0.0.1",
 					},
-					Interval: uint64(100),
+					Interval: uint64(60),
 				},
 			},
 			wantResp: &c2pb.ClaimTasksResponse{},
 			wantCode: codes.OK,
+
+			wantBeaconExist:            true,
+			wantBeaconLastSeenAtBefore: time.Now().UTC().Add(10 * time.Second),
+			wantBeaconLastSeenAtAfter:  time.Now().UTC().Add(-10 * time.Second),
 		},
 		{
 			name: "Second_Callback",
@@ -76,9 +86,13 @@ func TestClaimTasks(t *testing.T) {
 			},
 			wantResp: &c2pb.ClaimTasksResponse{},
 			wantCode: codes.OK,
+
+			wantBeaconExist:            true,
+			wantBeaconLastSeenAtBefore: time.Now().UTC().Add(10 * time.Second),
+			wantBeaconLastSeenAtAfter:  time.Now().UTC().Add(-10 * time.Second),
 		},
 		{
-			name: "Callback_With_Tasks",
+			name: "Existing_Beacon",
 			req: &c2pb.ClaimTasksRequest{
 				Beacon: &c2pb.Beacon{
 					Identifier: existingBeacon.Identifier,
@@ -102,26 +116,46 @@ func TestClaimTasks(t *testing.T) {
 				},
 			},
 			wantCode: codes.OK,
+
+			wantBeaconExist:            true,
+			wantBeaconLastSeenAtBefore: time.Now().UTC().Add(10 * time.Second),
+			wantBeaconLastSeenAtAfter:  time.Now().UTC().Add(-10 * time.Second),
 		},
 	}
 
-	testHandler := func(t *testing.T, tc testCase) {
-		resp, err := client.ClaimTasks(ctx, tc.req)
-		require.Equal(t, tc.wantCode.String(), status.Code(err).String(), err)
-		if status.Code(err) != codes.OK {
-			// Do not continue if we expected error code
-			return
-		}
+	// Run Tests
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			// Callback
+			resp, err := client.ClaimTasks(ctx, tc.req)
 
-		if diff := cmp.Diff(tc.wantResp, resp, protocmp.Transform()); diff != "" {
-			t.Errorf("invalid response (-want +got): %v", diff)
-		}
+			// Assert Response Code
+			require.Equal(t, tc.wantCode.String(), status.Code(err).String(), err)
+			if status.Code(err) != codes.OK {
+				// Do not continue if we expected error code
+				return
+			}
 
-	}
+			// Assert Response
+			if diff := cmp.Diff(tc.wantResp, resp, protocmp.Transform()); diff != "" {
+				t.Errorf("invalid response (-want +got): %v", diff)
+			}
 
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			testHandler(t, tc)
+			// Load Beacon
+			testBeacon, err := graph.Beacon.Query().
+				Where(
+					beacon.Identifier(tc.req.Beacon.Identifier),
+				).Only(ctx)
+			if ent.IsNotFound(err) && !tc.wantBeaconExist {
+				return
+			}
+			if err != nil {
+				t.Errorf("failed to load beacon: %v", err)
+				return
+			}
+
+			// Beacon Assertions
+			assert.WithinRange(t, testBeacon.LastSeenAt, tc.wantBeaconLastSeenAtAfter, tc.wantBeaconLastSeenAtBefore)
 		})
 	}
 }

tavern/internal/ent/schema/beacon.go

@@ -6,13 +6,12 @@ import (
 	"fmt"
 	"io"
 
-	"realm.pub/tavern/internal/namegen"
-
 	"entgo.io/contrib/entgql"
 	"entgo.io/ent"
 	"entgo.io/ent/schema"
 	"entgo.io/ent/schema/edge"
 	"entgo.io/ent/schema/field"
+	"realm.pub/tavern/internal/namegen"
 )
 
 // Beacon holds the schema definition for the Beacon entity.