github.com/david-imola/snapd@v0.0.0-20210611180407-2de8ddeece6d/overlord/servicestate/quotas_test.go (about)

     1  // -*- Mode: Go; indent-tabs-mode: t -*-
     2  
     3  /*
     4   * Copyright (C) 2021 Canonical Ltd
     5   *
     6   * This program is free software: you can redistribute it and/or modify
     7   * it under the terms of the GNU General Public License version 3 as
     8   * published by the Free Software Foundation.
     9   *
    10   * This program is distributed in the hope that it will be useful,
    11   * but WITHOUT ANY WARRANTY; without even the implied warranty of
    12   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    13   * GNU General Public License for more details.
    14   *
    15   * You should have received a copy of the GNU General Public License
    16   * along with this program.  If not, see <http://www.gnu.org/licenses/>.
    17   *
    18   */
    19  
    20  package servicestate_test
    21  
    22  import (
    23  	. "gopkg.in/check.v1"
    24  
    25  	"github.com/snapcore/snapd/gadget/quantity"
    26  	"github.com/snapcore/snapd/overlord/servicestate"
    27  	"github.com/snapcore/snapd/snap/quota"
    28  )
    29  
    30  type servicestateQuotasSuite struct {
    31  	baseServiceMgrTestSuite
    32  }
    33  
    34  var _ = Suite(&servicestateQuotasSuite{})
    35  
    36  func (s *servicestateQuotasSuite) SetUpTest(c *C) {
    37  	s.baseServiceMgrTestSuite.SetUpTest(c)
    38  
    39  	// we don't need the EnsureSnapServices ensure loop to run by default
    40  	servicestate.MockEnsuredSnapServices(s.mgr, true)
    41  }
    42  
    43  func (s *servicestateQuotasSuite) TestQuotas(c *C) {
    44  	st := s.state
    45  	st.Lock()
    46  	defer st.Unlock()
    47  
    48  	// with nothing in state we don't get any quotas
    49  	quotaMap, err := servicestate.AllQuotas(st)
    50  	c.Assert(err, IsNil)
    51  	c.Assert(quotaMap, HasLen, 0)
    52  
    53  	// we can add some basic quotas to state
    54  	grp := &quota.Group{
    55  		Name:        "foogroup",
    56  		MemoryLimit: quantity.SizeGiB,
    57  	}
    58  	newGrps, err := servicestate.PatchQuotas(st, grp)
    59  	c.Assert(err, IsNil)
    60  	c.Assert(newGrps, DeepEquals, map[string]*quota.Group{
    61  		"foogroup": grp,
    62  	})
    63  
    64  	// now we get back the same quota
    65  	quotaMap, err = servicestate.AllQuotas(st)
    66  	c.Assert(err, IsNil)
    67  	c.Assert(quotaMap, DeepEquals, map[string]*quota.Group{
    68  		"foogroup": grp,
    69  	})
    70  
    71  	// adding a sub-group quota only works when we update the parent group to
    72  	// reference the sub-group at the same time
    73  	grp2 := &quota.Group{
    74  		Name:        "group-2",
    75  		MemoryLimit: quantity.SizeGiB,
    76  		ParentGroup: "foogroup",
    77  	}
    78  	_, err = servicestate.PatchQuotas(st, grp2)
    79  	c.Assert(err, ErrorMatches, `cannot update quota "group-2": group "foogroup" does not reference necessary child group "group-2"`)
    80  
    81  	// we also can't add a sub-group to the parent without adding the sub-group
    82  	// itself
    83  	grp.SubGroups = append(grp.SubGroups, "group-2")
    84  	_, err = servicestate.PatchQuotas(st, grp)
    85  	c.Assert(err, ErrorMatches, `cannot update quota "foogroup": missing group "group-2" referenced as the sub-group of group "foogroup"`)
    86  
    87  	// foogroup didn't get updated in the state to mention the sub-group
    88  	foogrp, err := servicestate.GetQuota(st, "foogroup")
    89  	c.Assert(err, IsNil)
    90  	c.Assert(foogrp.SubGroups, HasLen, 0)
    91  
    92  	// but if we update them both at the same time we succeed
    93  	newGrps, err = servicestate.PatchQuotas(st, grp, grp2)
    94  	c.Assert(err, IsNil)
    95  	c.Assert(newGrps, DeepEquals, map[string]*quota.Group{
    96  		"foogroup": grp,
    97  		"group-2":  grp2,
    98  	})
    99  
   100  	// and now we see both in the state
   101  	quotaMap, err = servicestate.AllQuotas(st)
   102  	c.Assert(err, IsNil)
   103  	c.Assert(quotaMap, DeepEquals, map[string]*quota.Group{
   104  		"foogroup": grp,
   105  		"group-2":  grp2,
   106  	})
   107  
   108  	// and we can get individual quotas too
   109  	res, err := servicestate.GetQuota(st, "foogroup")
   110  	c.Assert(err, IsNil)
   111  	c.Assert(res, DeepEquals, grp)
   112  
   113  	res2, err := servicestate.GetQuota(st, "group-2")
   114  	c.Assert(err, IsNil)
   115  	c.Assert(res2, DeepEquals, grp2)
   116  
   117  	_, err = servicestate.GetQuota(st, "unknown")
   118  	c.Assert(err, Equals, servicestate.ErrQuotaNotFound)
   119  
   120  	// adding multiple quotas that are invalid produces a nice error message
   121  	otherGrp := &quota.Group{
   122  		Name: "other-group",
   123  		// invalid memory limit
   124  	}
   125  
   126  	otherGrp2 := &quota.Group{
   127  		Name: "other-group2",
   128  		// invalid memory limit
   129  	}
   130  
   131  	_, err = servicestate.PatchQuotas(st, otherGrp2, otherGrp)
   132  	// either group can get checked first
   133  	c.Assert(err, ErrorMatches, `cannot update quotas "other-group", "other-group2": group "other-group2?" is invalid: group memory limit must be non-zero`)
   134  }