Method TestAIBridge

coderd/database/dbauthz/dbauthz_test.go:6186–6615

()

}

func (s *MethodTestSuite) TestAIBridge() {
	s.Run("InsertAIBridgeInterception", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
		initID := uuid.UUID{3}
		user := testutil.Fake(s.T(), faker, database.User{ID: initID})
		// testutil.Fake cannot distinguish between a zero value and an explicitly requested value which is equivalent.
		user.IsSystem = false
		user.Deleted = false

		intID := uuid.UUID{2}
		intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID, InitiatorID: initID})

		params := database.InsertAIBridgeInterceptionParams{ID: intc.ID, InitiatorID: intc.InitiatorID, Provider: intc.Provider, Model: intc.Model}
		db.EXPECT().GetUserByID(gomock.Any(), initID).Return(user, nil).AnyTimes() // Validation.
		db.EXPECT().InsertAIBridgeInterception(gomock.Any(), params).Return(intc, nil).AnyTimes()
		check.Args(params).Asserts(intc, policy.ActionCreate)
	}))

	s.Run("InsertAIBridgeModelThought", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
		intID := uuid.UUID{2}
		intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID})
		db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation.

		params := database.InsertAIBridgeModelThoughtParams{InterceptionID: intc.ID}
		expected := testutil.Fake(s.T(), faker, database.AIBridgeModelThought{InterceptionID: intc.ID})
		db.EXPECT().InsertAIBridgeModelThought(gomock.Any(), params).Return(expected, nil).AnyTimes()
		check.Args(params).Asserts(intc, policy.ActionUpdate)
	}))

	s.Run("InsertAIBridgeTokenUsage", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
		intID := uuid.UUID{2}
		intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID})
		db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation.

		params := database.InsertAIBridgeTokenUsageParams{InterceptionID: intc.ID}
		expected := testutil.Fake(s.T(), faker, database.AIBridgeTokenUsage{InterceptionID: intc.ID})
		db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), params).Return(expected, nil).AnyTimes()
		check.Args(params).Asserts(intc, policy.ActionUpdate)
	}))

	s.Run("InsertAIBridgeUserPrompt", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
		intID := uuid.UUID{2}
		intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID})
		db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation.

		params := database.InsertAIBridgeUserPromptParams{InterceptionID: intc.ID}
		expected := testutil.Fake(s.T(), faker, database.AIBridgeUserPrompt{InterceptionID: intc.ID})
		db.EXPECT().InsertAIBridgeUserPrompt(gomock.Any(), params).Return(expected, nil).AnyTimes()
		check.Args(params).Asserts(intc, policy.ActionUpdate)
	}))

	s.Run("InsertAIBridgeToolUsage", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
		intID := uuid.UUID{2}
		intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID})
		db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation.

		params := database.InsertAIBridgeToolUsageParams{InterceptionID: intc.ID}
		expected := testutil.Fake(s.T(), faker, database.AIBridgeToolUsage{InterceptionID: intc.ID})
		db.EXPECT().InsertAIBridgeToolUsage(gomock.Any(), params).Return(expected, nil).AnyTimes()