Merge pull request juju#18054 from manadart/dqlite-unit-state-service

juju#18054

This adds a service for setting unit state, that recruits the persistence added under juju#18046.

Each of the state parts are set selectively based on what is populated in the `AgentState` struct. There are two call sites that set this state, and they populate different subsets of the state.

Those sites will recruit this service in a patch to follow.

## QA steps

Service is not yet recruited. Unit tests verify.

## Documentation changes

None.

## Links

**Jira card:** [JUJU-6026](https://warthogs.atlassian.net/browse/JUJU-6026)



[JUJU-6026]: https://warthogs.atlassian.net/browse/JUJU-6026?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

domain/unitstate/service/package_test.go

@@ -0,0 +1,19 @@
+// Copyright 2024 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package service
+
+import (
+	"testing"
+
+	gc "gopkg.in/check.v1"
+)
+
+//go:generate go run go.uber.org/mock/mockgen -typed -package service -destination state_mock_test.go github.com/juju/juju/domain/unitstate/service State
+func TestPackage(t *testing.T) {
+	gc.TestingT(t)
+}
+
+func ptr[T any](v T) *T {
+	return &v
+}

domain/unitstate/service/service.go

@@ -0,0 +1,104 @@
+// Copyright 2024 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package service
+
+import (
+	"context"
+
+	"github.com/juju/juju/domain"
+	"github.com/juju/juju/domain/unitstate"
+	"github.com/juju/juju/internal/errors"
+)
+
+// State defines an interface for interacting with the underlying state.
+type State interface {
+	domain.AtomicStateBase
+
+	// GetUnitUUIDForName returns the UUID for
+	// the unit identified by the input name.
+	GetUnitUUIDForName(domain.AtomicContext, string) (string, error)
+
+	// EnsureUnitStateRecord ensures that there is a record
+	// for the agent state for the unit with the input UUID.
+	EnsureUnitStateRecord(domain.AtomicContext, string) error
+
+	// UpdateUnitStateUniter updates the agent uniter
+	// state for the unit with the input UUID.
+	UpdateUnitStateUniter(domain.AtomicContext, string, string) error
+
+	// UpdateUnitStateStorage updates the agent storage
+	// state for the unit with the input UUID.
+	UpdateUnitStateStorage(domain.AtomicContext, string, string) error
+
+	// UpdateUnitStateSecret updates the agent secret
+	// state for the unit with the input UUID.
+	UpdateUnitStateSecret(domain.AtomicContext, string, string) error
+
+	// UpdateUnitStateCharm updates the agent charm
+	// state for the unit with the input UUID.
+	UpdateUnitStateCharm(domain.AtomicContext, string, map[string]string) error
+
+	// UpdateUnitStateRelation updates the agent relation
+	// state for the unit with the input UUID.
+	UpdateUnitStateRelation(domain.AtomicContext, string, map[int]string) error
+}
+
+// Service defines a service for interacting with the underlying state.
+type Service struct {
+	st State
+}
+
+// NewService returns a new Service for interacting with the underlying state.
+func NewService(st State) *Service {
+	return &Service{
+		st: st,
+	}
+}
+
+// SetState persists the input agent state selectively,
+// based on its populated values.
+func (s *Service) SetState(ctx context.Context, as unitstate.AgentState) error {
+	return s.st.RunAtomic(ctx, func(ctx domain.AtomicContext) error {
+		uuid, err := s.st.GetUnitUUIDForName(ctx, as.Name)
+		if err != nil {
+			return errors.Errorf("getting unit UUID for %s: %w", as.Name, err)
+		}
+
+		if err = s.st.EnsureUnitStateRecord(ctx, uuid); err != nil {
+			return errors.Errorf("ensuring state record for %s: %w", as.Name, err)
+		}
+
+		if as.UniterState != nil {
+			if err = s.st.UpdateUnitStateUniter(ctx, uuid, *as.UniterState); err != nil {
+				return errors.Errorf("setting uniter state for %s: %w", as.Name, err)
+			}
+		}
+
+		if as.StorageState != nil {
+			if err = s.st.UpdateUnitStateStorage(ctx, uuid, *as.StorageState); err != nil {
+				return errors.Errorf("setting storage state for %s: %w", as.Name, err)
+			}
+		}
+
+		if as.SecretState != nil {
+			if err = s.st.UpdateUnitStateSecret(ctx, uuid, *as.SecretState); err != nil {
+				return errors.Errorf("setting secret state for %s: %w", as.Name, err)
+			}
+		}
+
+		if as.CharmState != nil {
+			if err = s.st.UpdateUnitStateCharm(ctx, uuid, *as.CharmState); err != nil {
+				return errors.Errorf("setting charm state for %s: %w", as.Name, err)
+			}
+		}
+
+		if as.RelationState != nil {
+			if err = s.st.UpdateUnitStateRelation(ctx, uuid, *as.RelationState); err != nil {
+				return errors.Errorf("setting relation state for %s: %w", as.Name, err)
+			}
+		}
+
+		return nil
+	})
+}

domain/unitstate/service/service_test.go

@@ -0,0 +1,90 @@
+// Copyright 2024 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package service
+
+import (
+	"context"
+
+	jc "github.com/juju/testing/checkers"
+	"go.uber.org/mock/gomock"
+	gc "gopkg.in/check.v1"
+
+	"github.com/juju/juju/domain"
+	"github.com/juju/juju/domain/application/errors"
+	domaintesting "github.com/juju/juju/domain/testing"
+	"github.com/juju/juju/domain/unitstate"
+	unitstateerrors "github.com/juju/juju/domain/unitstate/errors"
+)
+
+type serviceSuite struct {
+	st *MockState
+}
+
+var _ = gc.Suite(&serviceSuite{})
+
+func (s *serviceSuite) TestSetStateAllAttributes(c *gc.C) {
+	defer s.setupMocks(c).Finish()
+
+	uuid := "some-unit-uuid"
+
+	exp := s.st.EXPECT()
+	exp.GetUnitUUIDForName(gomock.Any(), "unit/0").Return(uuid, nil)
+	exp.EnsureUnitStateRecord(gomock.Any(), uuid).Return(nil)
+	exp.UpdateUnitStateUniter(gomock.Any(), uuid, "some-uniter-state-yaml").Return(nil)
+	exp.UpdateUnitStateStorage(gomock.Any(), uuid, "some-storage-state-yaml").Return(nil)
+	exp.UpdateUnitStateSecret(gomock.Any(), uuid, "some-secret-state-yaml").Return(nil)
+	exp.UpdateUnitStateCharm(gomock.Any(), uuid, map[string]string{"one-key": "one-value"}).Return(nil)
+	exp.UpdateUnitStateRelation(gomock.Any(), uuid, map[int]string{1: "one-value"}).Return(nil)
+
+	err := NewService(s.st).SetState(context.Background(), unitstate.AgentState{
+		Name:          "unit/0",
+		CharmState:    ptr(map[string]string{"one-key": "one-value"}),
+		UniterState:   ptr("some-uniter-state-yaml"),
+		RelationState: ptr(map[int]string{1: "one-value"}),
+		StorageState:  ptr("some-storage-state-yaml"),
+		SecretState:   ptr("some-secret-state-yaml"),
+	})
+	c.Assert(err, jc.ErrorIsNil)
+}
+
+func (s *serviceSuite) TestSetStateSubsetAttributes(c *gc.C) {
+	defer s.setupMocks(c).Finish()
+
+	uuid := "some-unit-uuid"
+
+	exp := s.st.EXPECT()
+	exp.GetUnitUUIDForName(gomock.Any(), "unit/0").Return(uuid, nil)
+	exp.EnsureUnitStateRecord(gomock.Any(), uuid).Return(nil)
+	exp.UpdateUnitStateUniter(gomock.Any(), uuid, "some-uniter-state-yaml").Return(nil)
+
+	err := NewService(s.st).SetState(context.Background(), unitstate.AgentState{
+		Name:        "unit/0",
+		UniterState: ptr("some-uniter-state-yaml"),
+	})
+	c.Assert(err, jc.ErrorIsNil)
+}
+
+func (s *serviceSuite) TestSetStateUnitNotFound(c *gc.C) {
+	defer s.setupMocks(c).Finish()
+
+	exp := s.st.EXPECT()
+	exp.GetUnitUUIDForName(gomock.Any(), "unit/0").Return("", errors.UnitNotFound)
+
+	err := NewService(s.st).SetState(context.Background(), unitstate.AgentState{
+		Name:        "unit/0",
+		UniterState: ptr("some-uniter-state-yaml"),
+	})
+	c.Check(err, jc.ErrorIs, unitstateerrors.UnitNotFound)
+}
+
+func (s *serviceSuite) setupMocks(c *gc.C) *gomock.Controller {
+	ctrl := gomock.NewController(c)
+
+	s.st = NewMockState(ctrl)
+	s.st.EXPECT().RunAtomic(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, fn func(ctx domain.AtomicContext) error) error {
+		return fn(domaintesting.NewAtomicContext(ctx))
+	}).AnyTimes()
+
+	return ctrl
+}