test(session): cover machine proof persistence

Co-Authored-By: Codex <noreply@openai.com>

Author: cardyok

cmd/gpud/run/mock_command_test.go

@@ -260,15 +260,22 @@ func TestGetSessionCredentialsOptions_AllCredentialsPresent(t *testing.T) {
 		}).Build()
 
 		mockey.Mock(pkgmetadata.ReadMetadata).To(func(ctx context.Context, db *sql.DB, key string) (string, error) {
-			if key == pkgmetadata.MetadataKeyEndpoint {
+			switch key {
+			case pkgmetadata.MetadataKeyEndpoint:
 				return "https://stored.endpoint.com", nil
+			case pkgmetadata.MetadataKeyMachineProof:
+				return "machine-proof", nil
 			}
 			return "", nil
 		}).Build()
 
 		opts := getSessionCredentialsOptions(true, tmpDir, "")
 		assert.NotEmpty(t, opts)
-		assert.Len(t, opts, 3) // token, machine ID, endpoint
+		assert.Len(t, opts, 4) // token, machine ID, endpoint, machine proof
+
+		op := &config.Op{}
+		require.NoError(t, op.ApplyOpts(opts))
+		assert.Equal(t, "machine-proof", op.SessionMachineProof)
 	})
 }
 

cmd/gpud/run/mock_session_test.go

@@ -38,13 +38,18 @@ func TestReadSessionCredentialsFromPersistentFile(t *testing.T) {
 	require.NoError(t, pkgmetadata.SetMetadata(ctx, dbRW, pkgmetadata.MetadataKeyToken, "session-token"))
 	require.NoError(t, pkgmetadata.SetMetadata(ctx, dbRW, pkgmetadata.MetadataKeyMachineID, "assigned-machine-id"))
 	require.NoError(t, pkgmetadata.SetMetadata(ctx, dbRW, pkgmetadata.MetadataKeyEndpoint, "gpud-manager.example.com"))
+	require.NoError(t, pkgmetadata.SetMetadata(ctx, dbRW, pkgmetadata.MetadataKeyMachineProof, "machine-proof"))
 
 	gotToken, gotMachineID, gotEndpoint, err := readSessionCredentialsFromPersistentFile(ctx, tmpDir)
 	require.NoError(t, err)
 
 	assert.Equal(t, "session-token", gotToken)
 	assert.Equal(t, "assigned-machine-id", gotMachineID)
 	assert.Equal(t, "gpud-manager.example.com", gotEndpoint)
+
+	gotMachineProof, err := readMachineProofFromPersistentFile(ctx, tmpDir)
+	require.NoError(t, err)
+	assert.Equal(t, "machine-proof", gotMachineProof)
 }
 
 func TestReadSessionCredentialsFromPersistentFile_StateFileNotFound(t *testing.T) {

pkg/config/options_test.go

@@ -147,6 +147,7 @@ func TestDBInMemoryWithSessionCredentials(t *testing.T) {
 		WithDBInMemory(true),
 		WithSessionToken("session-token-from-login-response"),
 		WithSessionMachineID("assigned-machine-id-from-login-response"),
+		WithSessionMachineProof("machine-proof-from-login-response"),
 		WithSessionEndpoint("https://api.example.com"),
 	}
 
@@ -156,6 +157,7 @@ func TestDBInMemoryWithSessionCredentials(t *testing.T) {
 	assert.True(t, op.DBInMemory)
 	assert.Equal(t, "session-token-from-login-response", op.SessionToken)
 	assert.Equal(t, "assigned-machine-id-from-login-response", op.SessionMachineID)
+	assert.Equal(t, "machine-proof-from-login-response", op.SessionMachineProof)
 	assert.Equal(t, "https://api.example.com", op.SessionEndpoint)
 }
 

pkg/server/server_test.go

@@ -94,10 +94,11 @@ func TestNew_DBInMemory_SeedsSessionCredentialsBeforeAddressValidation(t *testin
 		MetricsRetentionPeriod: metav1.Duration{Duration: time.Minute},
 		Components:             []string{"-disable-all"},
 
-		DBInMemory:       true,
-		SessionToken:     "session-token",
-		SessionMachineID: "assigned-machine-id",
-		SessionEndpoint:  "https://api.example.com",
+		DBInMemory:          true,
+		SessionToken:        "session-token",
+		SessionMachineID:    "assigned-machine-id",
+		SessionMachineProof: "machine-proof",
+		SessionEndpoint:     "https://api.example.com",
 	}
 
 	s, err := New(ctx, log.NewNopAuditLogger(), cfg, nil)
@@ -106,6 +107,26 @@ func TestNew_DBInMemory_SeedsSessionCredentialsBeforeAddressValidation(t *testin
 	assert.Contains(t, err.Error(), "failed to create local GPUd server endpoint")
 }
 
+func TestReadMachineProof(t *testing.T) {
+	ctx := context.Background()
+	assert.Empty(t, (&Server{}).readMachineProof(ctx))
+
+	stateFile := filepath.Join(t.TempDir(), "state.db")
+	dbRW, err := sqlite.Open(stateFile)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = dbRW.Close() })
+	require.NoError(t, pkgmetadata.CreateTableMetadata(ctx, dbRW))
+	require.NoError(t, pkgmetadata.SetMetadata(ctx, dbRW, pkgmetadata.MetadataKeyMachineProof, "machine-proof"))
+
+	dbRO, err := sqlite.Open(stateFile, sqlite.WithReadOnly(true))
+	require.NoError(t, err)
+	s := &Server{dbRO: dbRO}
+	assert.Equal(t, "machine-proof", s.readMachineProof(ctx))
+
+	require.NoError(t, dbRO.Close())
+	assert.Empty(t, s.readMachineProof(ctx))
+}
+
 func TestGenerateSelfSignedCert(t *testing.T) {
 	s := &Server{}
 	cert, err := s.generateSelfSignedCert()

pkg/session/session_test.go

@@ -82,8 +82,9 @@ func TestNewSession(t *testing.T) {
 
 	endpoint := "test-endpoint.com"
 	machineID := "test-machine-id"
+	machineProof := "test-machine-proof"
 
-	session, err := NewSession(ctx, "", endpoint, "", WithMachineID(machineID), WithPipeInterval(time.Second), WithEnableAutoUpdate(true), WithComponentsRegistry(components.NewRegistry(nil)))
+	session, err := NewSession(ctx, "", endpoint, "", WithMachineID(machineID), WithMachineProof(machineProof), WithPipeInterval(time.Second), WithEnableAutoUpdate(true), WithComponentsRegistry(components.NewRegistry(nil)))
 	if err != nil {
 		t.Fatalf("error creating session: %v", err)
 	}
@@ -98,6 +99,9 @@ func TestNewSession(t *testing.T) {
 	if session.machineID != machineID {
 		t.Errorf("expected machineID %s, got %s", machineID, session.machineID)
 	}
+	if session.machineProof != machineProof {
+		t.Errorf("expected machineProof %s, got %s", machineProof, session.machineProof)
+	}
 }
 
 func TestCreateSessionRequestIncludesMachineProof(t *testing.T) {