def testAddLabels(self):
client_ref = self.api.Client(client_id=self.client_id)
self.assertEqual(list(client_ref.Get().data.labels), [])
with test_lib.FakeTime(42):
client_ref.AddLabels(["foo", "bar"])
self.assertCountEqual(client_ref.Get().data.labels, [
objects_pb2.ClientLabel(name="bar", owner=self.test_username),
objects_pb2.ClientLabel(name="foo", owner=self.test_username)
])
def testAddLabels(self):
client_ref = self.api.Client(client_id=self.client_urn.Basename())
self.assertEqual(list(client_ref.Get().data.labels), [])
with test_lib.FakeTime(42):
client_ref.AddLabels(["foo", "bar"])
self.assertEqual(
sorted(client_ref.Get().data.labels, key=lambda l: l.name), [
objects_pb2.ClientLabel(name="bar", owner=self.token.username),
objects_pb2.ClientLabel(name="foo", owner=self.token.username)
])
def testRemoveLabel(self):
with test_lib.FakeTime(42):
data_store.REL_DB.AddClientLabels(self.client_id, self.test_username,
["bar", "foo"])
client_ref = self.api.Client(client_id=self.client_id)
self.assertCountEqual(client_ref.Get().data.labels, [
objects_pb2.ClientLabel(name="bar", owner=self.test_username),
objects_pb2.ClientLabel(name="foo", owner=self.test_username)
])
client_ref.RemoveLabel("foo")
self.assertCountEqual(
client_ref.Get().data.labels,
[objects_pb2.ClientLabel(name="bar", owner=self.test_username)])
def testAddLabelsWithGeneratorArg(self):
client_ref = self.api.Client(client_id=self.client_id)
self.assertEqual(list(client_ref.Get().data.labels), [])
def Gen():
yield "foo"
yield "bar"
with test_lib.FakeTime(42):
client_ref.AddLabels(Gen())
self.assertCountEqual(client_ref.Get().data.labels, [
objects_pb2.ClientLabel(name="bar", owner=self.test_username),
objects_pb2.ClientLabel(name="foo", owner=self.test_username)
])
def testRemoveLabels(self):
with test_lib.FakeTime(42):
with aff4.FACTORY.Open(self.client_urn,
aff4_type=aff4_grr.VFSGRRClient,
mode="rw",
token=self.token) as client_obj:
client_obj.AddLabels(["bar", "foo"])
client_ref = self.api.Client(client_id=self.client_urn.Basename())
self.assertEqual(
sorted(client_ref.Get().data.labels, key=lambda l: l.name), [
objects_pb2.ClientLabel(name="bar", owner=self.token.username),
objects_pb2.ClientLabel(name="foo", owner=self.token.username)
])
client_ref.RemoveLabel("foo")
self.assertEqual(
sorted(client_ref.Get().data.labels, key=lambda l: l.name),
[objects_pb2.ClientLabel(name="bar", owner=self.token.username)])
