#!/usr/bin/env python
"""Tests for the hunt database api."""

import collections
import random
from typing import List, Optional

from google.protobuf import any_pb2
from grr_response_core.lib import rdfvalue
from grr_response_core.lib.rdfvalues import client as rdf_client
from grr_response_core.lib.rdfvalues import client_stats as rdf_client_stats
from grr_response_core.lib.rdfvalues import structs as rdf_structs
from grr_response_proto import flows_pb2
from grr_response_proto import hunts_pb2
from grr_response_proto import jobs_pb2
from grr_response_proto import objects_pb2
from grr_response_proto import output_plugin_pb2
from grr_response_server import flow
from grr_response_server.databases import db
from grr_response_server.databases import db_test_utils
from grr_response_server.models import hunts as model_hunts
from grr_response_server.output_plugins import email_plugin
from grr_response_server.rdfvalues import flow_objects as rdf_flow_objects
from grr_response_server.rdfvalues import hunt_objects as rdf_hunt_objects
from grr_response_server.rdfvalues import mig_flow_objects
from grr_response_server.rdfvalues import mig_hunt_objects
from grr_response_server.rdfvalues import mig_objects
from grr_response_server.rdfvalues import objects as rdf_objects


class DatabaseTestHuntMixin(object):
  """An abstract class for testing db.Database implementations.

  This mixin adds methods to test the handling of hunts.
  """

  def _SetupHuntClientAndFlow(
      self,
      hunt_id: str = None,
      client_id: Optional[str] = None,
      flow_id: Optional[str] = None,
      flow_state: Optional[rdf_structs.EnumNamedValue] = None,
      parent_flow_id: Optional[str] = None,
  ):
    client_id = db_test_utils.InitializeClient(self.db, client_id=client_id)
    # Top-level hunt-induced flows should have hunt's id.
    flow_id = flow_id or hunt_id

    flow_id = db_test_utils.InitializeFlow(
        self.db,
        client_id,
        flow_id=flow_id,
        flow_state=flow_state,
        parent_hunt_id=hunt_id,
        parent_flow_id=parent_flow_id,
    )

    return client_id, flow_id

  def testWritingAndReadingHuntObjectWorks(self):
    then = rdfvalue.RDFDatetime.Now()

    db_test_utils.InitializeUser(self.db, "Foo")
    hunt_id = db_test_utils.InitializeHunt(
        self.db,
        creator="Foo",
        description="Lorem ipsum.",
    )

    read_hunt_obj = self.db.ReadHuntObject(hunt_id)
    self.assertEqual(read_hunt_obj.creator, "Foo")
    self.assertEqual(read_hunt_obj.description, "Lorem ipsum.")
    self.assertGreater(
        rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
            read_hunt_obj.create_time
        ),
        then,
    )
    self.assertGreater(
        rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
            read_hunt_obj.last_update_time
        ),
        then,
    )

  def testWritingHuntObjectIntegralClientRate(self):
    creator = db_test_utils.InitializeUser(self.db, "user")

    hunt_id = db_test_utils.InitializeHunt(
        self.db,
        creator="user",
        description="Lorem ipsum.",
        client_rate=42,
    )

    hunt_obj = self.db.ReadHuntObject(hunt_id)
    self.assertEqual(hunt_obj.creator, creator)
    self.assertEqual(hunt_obj.description, "Lorem ipsum.")
    self.assertAlmostEqual(hunt_obj.client_rate, 42, places=5)

  def testWritingHuntObjectFractionalClientRate(self):
    creator = db_test_utils.InitializeUser(self.db, "user")

    hunt_id = db_test_utils.InitializeHunt(
        self.db,
        creator="user",
        description="Lorem ipsum.",
        client_rate=3.14,
    )

    hunt_obj = self.db.ReadHuntObject(hunt_id)
    self.assertEqual(hunt_obj.creator, creator)
    self.assertEqual(hunt_obj.description, "Lorem ipsum.")
    self.assertAlmostEqual(hunt_obj.client_rate, 3.14, places=5)

  def testHuntObjectCannotBeOverwritten(self):
    db_test_utils.InitializeUser(self.db, "user")
    hunt_id = "ABCDEF42"
    hunt_obj_v1 = hunts_pb2.Hunt(
        hunt_id=hunt_id,
        hunt_state=hunts_pb2.Hunt.HuntState.PAUSED,
        description="foo",
        creator="user",
    )
    hunt_obj_v2 = hunts_pb2.Hunt(
        hunt_id=hunt_id,
        hunt_state=hunts_pb2.Hunt.HuntState.PAUSED,
        description="bar",
        creator="user",
    )
    hunt_obj_v2.hunt_id = hunt_obj_v1.hunt_id

    self.db.WriteHuntObject(hunt_obj_v1)

    with self.assertRaises(db.DuplicatedHuntError) as context:
      self.db.WriteHuntObject(hunt_obj_v2)

    self.assertEqual(context.exception.hunt_id, hunt_id)

  def testHuntObjectCannotBeWrittenInNonPausedState(self):
    db_test_utils.InitializeUser(self.db, "user")
    hunt_object = hunts_pb2.Hunt(
        hunt_id=rdf_hunt_objects.RandomHuntId(),
        hunt_state=rdf_hunt_objects.Hunt.HuntState.STARTED,
        creator="user",
    )

    with self.assertRaises(ValueError):
      self.db.WriteHuntObject(hunt_object)

  def testReadingNonExistentHuntObjectRaises(self):
    with self.assertRaises(db.UnknownHuntError):
      self.db.ReadHuntObject(rdf_hunt_objects.RandomHuntId())

  def testUpdateHuntObjectRaisesIfHuntDoesNotExist(self):
    with self.assertRaises(db.UnknownHuntError):
      self.db.UpdateHuntObject(
          rdf_hunt_objects.RandomHuntId(),
          hunt_state=hunts_pb2.Hunt.HuntState.STARTED,
      )

  def testUpdateHuntObjectCorrectlyUpdatesHuntObject(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    self.db.UpdateHuntObject(
        hunt_id,
        duration=rdfvalue.Duration.From(1, rdfvalue.WEEKS),
        client_rate=33,
        client_limit=48,
        hunt_state=hunts_pb2.Hunt.HuntState.STOPPED,
        hunt_state_comment="foo",
        start_time=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(43),
        num_clients_at_start_time=44,
    )

    updated_hunt_obj = self.db.ReadHuntObject(hunt_id)
    self.assertEqual(
        rdfvalue.Duration.From(updated_hunt_obj.duration, rdfvalue.SECONDS),
        rdfvalue.Duration.From(1, rdfvalue.WEEKS),
    )
    self.assertEqual(updated_hunt_obj.client_rate, 33)
    self.assertEqual(updated_hunt_obj.client_limit, 48)
    self.assertEqual(
        updated_hunt_obj.hunt_state, hunts_pb2.Hunt.HuntState.STOPPED
    )
    self.assertEqual(updated_hunt_obj.hunt_state_comment, "foo")
    self.assertEqual(
        rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
            updated_hunt_obj.init_start_time
        ),
        rdfvalue.RDFDatetime.FromSecondsSinceEpoch(43),
    )
    self.assertEqual(
        rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
            updated_hunt_obj.last_start_time
        ),
        rdfvalue.RDFDatetime.FromSecondsSinceEpoch(43),
    )
    self.assertEqual(updated_hunt_obj.num_clients_at_start_time, 44)

  def testUpdateHuntObjectCorrectlyUpdatesInitAndLastStartTime(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    timestamp_1 = rdfvalue.RDFDatetime.Now()
    self.db.UpdateHuntObject(hunt_id, start_time=timestamp_1)

    timestamp_2 = rdfvalue.RDFDatetime.Now()
    self.db.UpdateHuntObject(hunt_id, start_time=timestamp_2)

    updated_hunt_object = self.db.ReadHuntObject(hunt_id)
    self.assertEqual(
        rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
            updated_hunt_object.init_start_time
        ),
        timestamp_1,
    )
    self.assertEqual(
        rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
            updated_hunt_object.last_start_time
        ),
        timestamp_2,
    )

  def testDeletingHuntObjectWorks(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    # This shouldn't raise.
    self.db.ReadHuntObject(hunt_id)

    self.db.DeleteHuntObject(hunt_id)

    # The hunt is deleted: this should raise now.
    with self.assertRaises(db.UnknownHuntError):
      self.db.ReadHuntObject(hunt_id)

  def testDeleteHuntObjectWithApprovalRequest(self):
    creator = db_test_utils.InitializeUser(self.db)
    approver = db_test_utils.InitializeUser(self.db)
    hunt_id = db_test_utils.InitializeHunt(self.db, creator=creator)

    approval = rdf_objects.ApprovalRequest()
    approval.approval_type = (
        rdf_objects.ApprovalRequest.ApprovalType.APPROVAL_TYPE_HUNT
    )
    approval.requestor_username = creator
    approval.notified_users = [approver]
    approval.subject_id = hunt_id
    approval.expiration_time = (
        rdfvalue.RDFDatetime.Now() + rdfvalue.Duration.From(1, rdfvalue.DAYS)
    )
    # TODO: Stop using `mig_objects`.
    proto_approval = mig_objects.ToProtoApprovalRequest(approval)
    approval_id = self.db.WriteApprovalRequest(proto_approval)

    self.db.DeleteHuntObject(hunt_id=hunt_id)

    with self.assertRaises(db.UnknownApprovalRequestError):
      self.db.ReadApprovalRequest(creator, approval_id)

  def testReadHuntObjectsReturnsEmptyListWhenNoHunts(self):
    self.assertEqual(self.db.ReadHuntObjects(offset=0, count=db.MAX_COUNT), [])

  def _CreateMultipleHunts(self) -> List[hunts_pb2.Hunt]:
    self._CreateMultipleUsers(["user-a", "user-b"])

    result = []
    for i in range(10):
      if i < 5:
        creator = "user-a"
      else:
        creator = "user-b"
      hunt_id = db_test_utils.InitializeHunt(
          self.db,
          creator=creator,
          description="foo_%d" % i,
      )
      result.append(self.db.ReadHuntObject(hunt_id))

    return result

  def _CreateMultipleUsers(self, users: List[str]):
    for user in users:
      db_test_utils.InitializeUser(self.db, user)

  def _CreateMultipleHuntsForUser(
      self,
      user: str,
      count: int,
  ) -> List[hunts_pb2.Hunt]:
    result = []
    for i in range(count):
      hunt_id = db_test_utils.InitializeHunt(
          self.db,
          creator=user,
          description="foo_%d" % i,
      )
      result.append(self.db.ReadHuntObject(hunt_id))
    return result

  def _CreateHuntWithState(
      self, creator: str, state: hunts_pb2.Hunt.HuntState
  ) -> hunts_pb2.Hunt:
    hunt_id = db_test_utils.InitializeHunt(self.db, creator=creator)

    self.db.UpdateHuntObject(hunt_id, hunt_state=state)
    return self.db.ReadHuntObject(hunt_id)

  def testReadHuntObjectsWithoutFiltersReadsAllHunts(self):
    expected = self._CreateMultipleHunts()
    got = self.db.ReadHuntObjects(0, db.MAX_COUNT)
    self.assertListEqual(got, list(reversed(expected)))

  def testReadHuntObjectsWithCreatorFilterIsAppliedCorrectly(self):
    all_hunts = self._CreateMultipleHunts()

    got = self.db.ReadHuntObjects(0, db.MAX_COUNT, with_creator="user-a")
    self.assertListEqual(got, list(reversed(all_hunts[:5])))

    got = self.db.ReadHuntObjects(0, db.MAX_COUNT, with_creator="user-b")
    self.assertListEqual(got, list(reversed(all_hunts[5:])))

  def testReadHuntObjectsWithCreatedByFilterIsAppliedCorrectly(self):
    self._CreateMultipleUsers(["user-a", "user-b"])
    all_hunts = self._CreateMultipleHuntsForUser(
        "user-a", 5
    ) + self._CreateMultipleHuntsForUser("user-b", 5)

    got = self.db.ReadHuntObjects(0, db.MAX_COUNT, created_by=frozenset([]))
    self.assertListEqual(got, [])

    got = self.db.ReadHuntObjects(
        0, db.MAX_COUNT, created_by=frozenset(["user-a"])
    )
    self.assertListEqual(got, list(reversed(all_hunts[:5])))

    got = self.db.ReadHuntObjects(
        0, db.MAX_COUNT, created_by=frozenset(["user-b"])
    )
    self.assertListEqual(got, list(reversed(all_hunts[5:])))

    got = self.db.ReadHuntObjects(
        0, db.MAX_COUNT, created_by=frozenset(["user-a", "user-b"])
    )
    self.assertListEqual(got, list(reversed(all_hunts)))

  def testReadHuntObjectsWithCreatorAndCreatedByFilterIsAppliedCorrectly(
      self,
  ):
    self._CreateMultipleUsers(["user-a", "user-b"])
    all_hunts = self._CreateMultipleHuntsForUser(
        "user-a", 5
    ) + self._CreateMultipleHuntsForUser("user-b", 5)

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-a",
        created_by=frozenset(["user-a"]),
    )
    self.assertListEqual(got, list(reversed(all_hunts[:5])))

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-b",
        created_by=frozenset(["user-b"]),
    )
    self.assertListEqual(got, list(reversed(all_hunts[5:])))

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-a",
        created_by=frozenset(["user-b"]),
    )
    self.assertListEqual(got, [])

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-b",
        created_by=frozenset(["user-a"]),
    )
    self.assertListEqual(got, [])

  def testReadHuntObjectsWithNotCreatedByFilterIsAppliedCorrectly(self):
    self._CreateMultipleUsers(["user-a", "user-b"])
    all_hunts = self._CreateMultipleHuntsForUser(
        "user-a", 5
    ) + self._CreateMultipleHuntsForUser("user-b", 5)

    got = self.db.ReadHuntObjects(0, db.MAX_COUNT, not_created_by=frozenset([]))
    self.assertListEqual(got, list(reversed(all_hunts)))

    got = self.db.ReadHuntObjects(
        0, db.MAX_COUNT, not_created_by=frozenset(["user-a"])
    )
    self.assertListEqual(got, list(reversed(all_hunts[5:])))

    got = self.db.ReadHuntObjects(
        0, db.MAX_COUNT, not_created_by=frozenset(["user-b"])
    )
    self.assertListEqual(got, list(reversed(all_hunts[:5])))

    got = self.db.ReadHuntObjects(
        0, db.MAX_COUNT, not_created_by=frozenset(["user-a", "user-b"])
    )
    self.assertListEqual(got, [])

  def testReadHuntObjectsWithCreatorAndNotCreatedByFilterIsAppliedCorrectly(
      self,
  ):
    self._CreateMultipleUsers(["user-a", "user-b"])
    all_hunts = self._CreateMultipleHuntsForUser(
        "user-a", 5
    ) + self._CreateMultipleHuntsForUser("user-b", 5)

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-a",
        not_created_by=frozenset([]),
    )
    self.assertListEqual(got, list(reversed(all_hunts[:5])))

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-a",
        not_created_by=frozenset(["user-a"]),
    )
    self.assertListEqual(got, [])

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-b",
        not_created_by=frozenset(["user-b"]),
    )
    self.assertListEqual(got, [])

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-a",
        not_created_by=frozenset(["user-b"]),
    )
    self.assertListEqual(got, list(reversed(all_hunts[:5])))

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-b",
        not_created_by=frozenset(["user-a"]),
    )
    self.assertListEqual(got, list(reversed(all_hunts[5:])))

  def testReadHuntObjectsCreatedAfterFilterIsAppliedCorrectly(self):
    all_hunts = self._CreateMultipleHunts()

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        created_after=rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
            all_hunts[0].create_time
        )
        - rdfvalue.Duration.From(1, rdfvalue.SECONDS),
    )
    self.assertListEqual(got, list(reversed(all_hunts)))

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        created_after=rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
            all_hunts[2].create_time
        ),
    )
    self.assertListEqual(got, list(reversed(all_hunts[3:])))

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        created_after=rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
            all_hunts[-1].create_time
        ),
    )
    self.assertEmpty(got)

  def testReadHuntObjectsWithDescriptionMatchFilterIsAppliedCorrectly(self):
    all_hunts = self._CreateMultipleHunts()

    got = self.db.ReadHuntObjects(
        0, db.MAX_COUNT, with_description_match="foo_"
    )
    self.assertListEqual(got, list(reversed(all_hunts)))

    got = self.db.ReadHuntObjects(
        0, db.MAX_COUNT, with_description_match="blah"
    )
    self.assertEmpty(got)

    got = self.db.ReadHuntObjects(
        0, db.MAX_COUNT, with_description_match="foo_3"
    )
    self.assertListEqual(got, [all_hunts[3]])

  def testReadHuntObjectsWithStatesFilterIsAppliedCorrectly(self):
    creator = "testuser"
    db_test_utils.InitializeUser(self.db, creator)
    paused_hunt = self._CreateHuntWithState(
        creator, hunts_pb2.Hunt.HuntState.PAUSED
    )
    self._CreateHuntWithState(creator, hunts_pb2.Hunt.HuntState.STARTED)
    stopped_hunt = self._CreateHuntWithState(
        creator, hunts_pb2.Hunt.HuntState.STOPPED
    )
    self._CreateHuntWithState(creator, hunts_pb2.Hunt.HuntState.COMPLETED)

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        with_states=[
            hunts_pb2.Hunt.HuntState.PAUSED,
        ],
    )
    self.assertListEqual(got, [paused_hunt])

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        with_states=[
            hunts_pb2.Hunt.HuntState.PAUSED,
            hunts_pb2.Hunt.HuntState.STOPPED,
        ],
    )
    self.assertCountEqual(got, [paused_hunt, stopped_hunt])

    got = self.db.ReadHuntObjects(
        0,
        db.MAX_COUNT,
        with_states=[],
    )
    self.assertListEqual(got, [])

  def testReadHuntObjectsCombinationsOfFiltersAreAppliedCorrectly(self):
    expected = self._CreateMultipleHunts()
    self.DoFilterCombinationsAndOffsetCountTest(
        self.db.ReadHuntObjects,
        conditions=dict(
            with_creator="user-a",
            created_after=rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
                expected[2].create_time
            ),
            with_description_match="foo_4",
        ),
        error_desc="ReadHuntObjects",
    )

  def testListHuntObjectsReturnsEmptyListWhenNoHunts(self):
    self.assertEqual(self.db.ListHuntObjects(offset=0, count=db.MAX_COUNT), [])

  def testListHuntObjectsWithoutFiltersReadsAllHunts(self):
    hunts = self._CreateMultipleHunts()
    hunts = [mig_hunt_objects.ToRDFHunt(h) for h in hunts]
    expected = [rdf_hunt_objects.HuntMetadata.FromHunt(h) for h in hunts]
    expected = [mig_hunt_objects.ToProtoHuntMetadata(h) for h in expected]
    got = self.db.ListHuntObjects(0, db.MAX_COUNT)
    self.assertListEqual(got, list(reversed(expected)))

  def testListHuntObjectsWithCreatorFilterIsAppliedCorrectly(self):
    all_hunts = self._CreateMultipleHunts()
    all_hunts = [mig_hunt_objects.ToRDFHunt(h) for h in all_hunts]
    all_hunts = [rdf_hunt_objects.HuntMetadata.FromHunt(h) for h in all_hunts]
    all_hunts = [mig_hunt_objects.ToProtoHuntMetadata(h) for h in all_hunts]

    got = self.db.ListHuntObjects(0, db.MAX_COUNT, with_creator="user-a")
    self.assertListEqual(got, list(reversed(all_hunts[:5])))

    got = self.db.ListHuntObjects(0, db.MAX_COUNT, with_creator="user-b")
    self.assertListEqual(got, list(reversed(all_hunts[5:])))

  def testListHuntObjectsWithCreatedByFilterIsAppliedCorrectly(self):
    self._CreateMultipleUsers(["user-a", "user-b"])
    hunts = self._CreateMultipleHuntsForUser(
        "user-a", 5
    ) + self._CreateMultipleHuntsForUser("user-b", 5)
    hunts = [mig_hunt_objects.ToRDFHunt(h) for h in hunts]
    all_hunts = [rdf_hunt_objects.HuntMetadata.FromHunt(h) for h in hunts]
    all_hunts = [mig_hunt_objects.ToProtoHuntMetadata(h) for h in all_hunts]

    got = self.db.ListHuntObjects(0, db.MAX_COUNT, created_by=frozenset([]))
    self.assertListEqual(got, [])

    got = self.db.ListHuntObjects(
        0, db.MAX_COUNT, created_by=frozenset(["user-a"])
    )
    self.assertListEqual(got, list(reversed(all_hunts[:5])))

    got = self.db.ListHuntObjects(
        0, db.MAX_COUNT, created_by=frozenset(["user-b"])
    )
    self.assertListEqual(got, list(reversed(all_hunts[5:])))

    got = self.db.ListHuntObjects(
        0, db.MAX_COUNT, created_by=frozenset(["user-a", "user-b"])
    )
    self.assertListEqual(got, list(reversed(all_hunts)))

  def testListHuntObjectsWithCreatorAndCreatedByFilterIsAppliedCorrectly(
      self,
  ):
    self._CreateMultipleUsers(["user-a", "user-b"])
    hunts = self._CreateMultipleHuntsForUser(
        "user-a", 5
    ) + self._CreateMultipleHuntsForUser("user-b", 5)
    hunts = [mig_hunt_objects.ToRDFHunt(h) for h in hunts]
    all_hunts = [rdf_hunt_objects.HuntMetadata.FromHunt(h) for h in hunts]
    all_hunts = [mig_hunt_objects.ToProtoHuntMetadata(h) for h in all_hunts]

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-a",
        created_by=frozenset(["user-a"]),
    )
    self.assertListEqual(got, list(reversed(all_hunts[:5])))

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-b",
        created_by=frozenset(["user-b"]),
    )
    self.assertListEqual(got, list(reversed(all_hunts[5:])))

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-a",
        created_by=frozenset(["user-b"]),
    )
    self.assertListEqual(got, [])

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-b",
        created_by=frozenset(["user-a"]),
    )
    self.assertListEqual(got, [])

  def testListHuntObjectsWithNotCreatedByFilterIsAppliedCorrectly(self):
    self._CreateMultipleUsers(["user-a", "user-b"])
    hunts = self._CreateMultipleHuntsForUser(
        "user-a", 5
    ) + self._CreateMultipleHuntsForUser("user-b", 5)
    hunts = [mig_hunt_objects.ToRDFHunt(h) for h in hunts]
    all_hunts = [rdf_hunt_objects.HuntMetadata.FromHunt(h) for h in hunts]
    all_hunts = [mig_hunt_objects.ToProtoHuntMetadata(h) for h in all_hunts]

    got = self.db.ListHuntObjects(0, db.MAX_COUNT, not_created_by=frozenset([]))
    self.assertListEqual(got, list(reversed(all_hunts)))

    got = self.db.ListHuntObjects(
        0, db.MAX_COUNT, not_created_by=frozenset(["user-a"])
    )
    self.assertListEqual(got, list(reversed(all_hunts[5:])))

    got = self.db.ListHuntObjects(
        0, db.MAX_COUNT, not_created_by=frozenset(["user-b"])
    )
    self.assertListEqual(got, list(reversed(all_hunts[:5])))

    got = self.db.ListHuntObjects(
        0, db.MAX_COUNT, not_created_by=frozenset(["user-a", "user-b"])
    )
    self.assertListEqual(got, [])

  def testListHuntObjectsWithCreatorAndNotCreatedByFilterIsAppliedCorrectly(
      self,
  ):
    self._CreateMultipleUsers(["user-a", "user-b"])
    hunts = self._CreateMultipleHuntsForUser(
        "user-a", 5
    ) + self._CreateMultipleHuntsForUser("user-b", 5)
    hunts = [mig_hunt_objects.ToRDFHunt(h) for h in hunts]

    all_hunts = [rdf_hunt_objects.HuntMetadata.FromHunt(h) for h in hunts]
    all_hunts = [mig_hunt_objects.ToProtoHuntMetadata(h) for h in all_hunts]

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-a",
        not_created_by=frozenset([]),
    )
    self.assertListEqual(got, list(reversed(all_hunts[:5])))

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-a",
        not_created_by=frozenset(["user-a"]),
    )
    self.assertListEqual(got, [])

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-b",
        not_created_by=frozenset(["user-b"]),
    )
    self.assertListEqual(got, [])

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-a",
        not_created_by=frozenset(["user-b"]),
    )
    self.assertListEqual(got, list(reversed(all_hunts[:5])))

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        with_creator="user-b",
        not_created_by=frozenset(["user-a"]),
    )
    self.assertListEqual(got, list(reversed(all_hunts[5:])))

  def testListHuntObjectsCreatedAfterFilterIsAppliedCorrectly(self):
    hunts = self._CreateMultipleHunts()
    hunts = [mig_hunt_objects.ToRDFHunt(h) for h in hunts]
    all_hunts = [rdf_hunt_objects.HuntMetadata.FromHunt(h) for h in hunts]
    all_hunts = [mig_hunt_objects.ToProtoHuntMetadata(h) for h in all_hunts]

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        created_after=rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
            all_hunts[0].create_time
        )
        - rdfvalue.Duration.From(1, rdfvalue.SECONDS),
    )
    self.assertListEqual(got, list(reversed(all_hunts)))

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        created_after=rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
            all_hunts[2].create_time
        ),
    )
    self.assertListEqual(got, list(reversed(all_hunts[3:])))

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        created_after=rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
            all_hunts[-1].create_time
        ),
    )
    self.assertEmpty(got)

  def testListHuntObjectsWithDescriptionMatchFilterIsAppliedCorrectly(self):
    hunts = self._CreateMultipleHunts()
    hunts = [mig_hunt_objects.ToRDFHunt(h) for h in hunts]
    all_hunts = [rdf_hunt_objects.HuntMetadata.FromHunt(h) for h in hunts]
    all_hunts = [mig_hunt_objects.ToProtoHuntMetadata(h) for h in all_hunts]

    got = self.db.ListHuntObjects(
        0, db.MAX_COUNT, with_description_match="foo_"
    )
    self.assertListEqual(got, list(reversed(all_hunts)))

    got = self.db.ListHuntObjects(
        0, db.MAX_COUNT, with_description_match="blah"
    )
    self.assertEmpty(got)

    got = self.db.ListHuntObjects(
        0, db.MAX_COUNT, with_description_match="foo_3"
    )
    self.assertListEqual(got, [all_hunts[3]])

  def testListHuntObjectsWithStatesFilterIsAppliedCorrectly(self):
    creator = "testuser"
    db_test_utils.InitializeUser(self.db, creator)
    paused_hunt = self._CreateHuntWithState(
        creator, hunts_pb2.Hunt.HuntState.PAUSED
    )
    paused_hunt = mig_hunt_objects.ToRDFHunt(paused_hunt)
    paused_hunt_metadata = rdf_hunt_objects.HuntMetadata.FromHunt(paused_hunt)
    paused_hunt_metadata = mig_hunt_objects.ToProtoHuntMetadata(
        paused_hunt_metadata
    )
    self._CreateHuntWithState(creator, rdf_hunt_objects.Hunt.HuntState.STARTED)
    stopped_hunt = self._CreateHuntWithState(
        creator, hunts_pb2.Hunt.HuntState.STOPPED
    )
    stopped_hunt = mig_hunt_objects.ToRDFHunt(stopped_hunt)
    stopped_hunt_metadata = rdf_hunt_objects.HuntMetadata.FromHunt(stopped_hunt)
    stopped_hunt_metadata = mig_hunt_objects.ToProtoHuntMetadata(
        stopped_hunt_metadata
    )
    self._CreateHuntWithState(creator, hunts_pb2.Hunt.HuntState.COMPLETED)

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        with_states=[
            hunts_pb2.Hunt.HuntState.PAUSED,
        ],
    )
    self.assertCountEqual(got, [paused_hunt_metadata])

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        with_states=[
            hunts_pb2.Hunt.HuntState.PAUSED,
            hunts_pb2.Hunt.HuntState.STOPPED,
        ],
    )
    self.assertCountEqual(got, [paused_hunt_metadata, stopped_hunt_metadata])

    got = self.db.ListHuntObjects(
        0,
        db.MAX_COUNT,
        with_states=[],
    )
    self.assertListEqual(got, [])

  def testListHuntObjectsCombinationsOfFiltersAreAppliedCorrectly(self):
    expected = self._CreateMultipleHunts()

    self.DoFilterCombinationsAndOffsetCountTest(
        self.db.ListHuntObjects,
        conditions=dict(
            with_creator="user-a",
            created_after=rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
                expected[2].create_time
            ),
            with_description_match="foo_4",
            created_by=frozenset(["user-a"]),
            not_created_by=frozenset(["user-b"]),
            with_states=[rdf_hunt_objects.Hunt.HuntState.PAUSED],
        ),
        error_desc="ListHuntObjects",
    )

  def testWritingAndReadingHuntOutputPluginsStatesWorks(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    email_args = output_plugin_pb2.EmailOutputPluginArgs(emails_limit=42)
    email_args_any = any_pb2.Any()
    email_args_any.Pack(email_args)
    plugin_descriptor = output_plugin_pb2.OutputPluginDescriptor(
        plugin_name=email_plugin.EmailOutputPlugin.__name__, args=email_args_any
    )
    plugin_state = jobs_pb2.AttributedDict(
        dat=[
            jobs_pb2.KeyValue(
                k=jobs_pb2.DataBlob(string="a_foo1"),
                v=jobs_pb2.DataBlob(string="a_bar1"),
            ),
            jobs_pb2.KeyValue(
                k=jobs_pb2.DataBlob(string="a_foo2"),
                v=jobs_pb2.DataBlob(string="a_bar2"),
            ),
        ]
    )
    state_1 = output_plugin_pb2.OutputPluginState(
        plugin_descriptor=plugin_descriptor,
        plugin_state=plugin_state,
    )

    email_args_2 = output_plugin_pb2.EmailOutputPluginArgs(emails_limit=43)
    email_args_any_2 = any_pb2.Any()
    email_args_any_2.Pack(email_args_2)
    plugin_descriptor_2 = output_plugin_pb2.OutputPluginDescriptor(
        plugin_name=email_plugin.EmailOutputPlugin.__name__,
        args=email_args_any_2,
    )
    plugin_state_2 = jobs_pb2.AttributedDict(
        dat=[
            jobs_pb2.KeyValue(
                k=jobs_pb2.DataBlob(string="b_foo1"),
                v=jobs_pb2.DataBlob(string="b_bar1"),
            ),
            jobs_pb2.KeyValue(
                k=jobs_pb2.DataBlob(string="b_foo2"),
                v=jobs_pb2.DataBlob(string="b_bar2"),
            ),
        ]
    )
    state_2 = output_plugin_pb2.OutputPluginState(
        plugin_descriptor=plugin_descriptor_2,
        plugin_state=plugin_state_2,
    )

    written_states = [state_1, state_2]
    self.db.WriteHuntOutputPluginsStates(hunt_id, written_states)

    read_states = self.db.ReadHuntOutputPluginsStates(hunt_id)
    self.assertEqual(read_states, written_states)

  def testReadingHuntOutputPluginsReturnsThemInOrderOfWriting(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    states = []
    for i in range(100):
      states.append(
          output_plugin_pb2.OutputPluginState(
              plugin_descriptor=output_plugin_pb2.OutputPluginDescriptor(
                  plugin_name="DummyHuntOutputPlugin_%d" % i
              ),
              plugin_state=jobs_pb2.AttributedDict(),
          )
      )
    random.shuffle(states)

    self.db.WriteHuntOutputPluginsStates(hunt_id, states)

    read_states = self.db.ReadHuntOutputPluginsStates(hunt_id)
    self.assertEqual(read_states, states)

  def testWritingHuntOutputStatesForZeroPlugins(self):
    # Passing an empty list of states is always a no-op so this should not
    # raise, even if the hunt does not exist.
    self.db.WriteHuntOutputPluginsStates(rdf_hunt_objects.RandomHuntId(), [])

  def testWritingHuntOutputStatesForUnknownHuntRaises(self):
    state = output_plugin_pb2.OutputPluginState(
        plugin_descriptor=output_plugin_pb2.OutputPluginDescriptor(
            plugin_name="DummyHuntOutputPlugin1"
        ),
        plugin_state=jobs_pb2.AttributedDict(),
    )

    with self.assertRaises(db.UnknownHuntError):
      self.db.WriteHuntOutputPluginsStates(
          rdf_hunt_objects.RandomHuntId(), [state]
      )

  def testReadingHuntOutputPluginsWithoutStates(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)
    res = self.db.ReadHuntOutputPluginsStates(hunt_id)
    self.assertEqual(res, [])

  def testReadingHuntOutputStatesForUnknownHuntRaises(self):
    with self.assertRaises(db.UnknownHuntError):
      self.db.ReadHuntOutputPluginsStates(rdf_hunt_objects.RandomHuntId())

  def testUpdatingHuntOutputStateForUnknownHuntRaises(self):
    with self.assertRaises(db.UnknownHuntError):
      self.db.UpdateHuntOutputPluginState(
          rdf_hunt_objects.RandomHuntId(), 0, lambda x: x
      )

  def testUpdatingHuntOutputStateWorksCorrectly(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    state_1 = output_plugin_pb2.OutputPluginState(
        plugin_descriptor=output_plugin_pb2.OutputPluginDescriptor(
            plugin_name="DummyHuntOutputPlugin1"
        ),
        plugin_state=jobs_pb2.AttributedDict(),
    )

    state_2 = output_plugin_pb2.OutputPluginState(
        plugin_descriptor=output_plugin_pb2.OutputPluginDescriptor(
            plugin_name="DummyHuntOutputPlugin2"
        ),
        plugin_state=jobs_pb2.AttributedDict(),
    )

    self.db.WriteHuntOutputPluginsStates(hunt_id, [state_1, state_2])

    def Update(s: jobs_pb2.AttributedDict) -> jobs_pb2.AttributedDict:
      el = s.dat.add()
      el.k.CopyFrom(jobs_pb2.DataBlob(string="foo"))
      el.v.CopyFrom(jobs_pb2.DataBlob(string="bar"))
      return s

    self.db.UpdateHuntOutputPluginState(hunt_id, 0, Update)

    states = self.db.ReadHuntOutputPluginsStates(hunt_id)
    self.assertEqual(
        states[0].plugin_state,
        jobs_pb2.AttributedDict(
            dat=[
                jobs_pb2.KeyValue(
                    k=jobs_pb2.DataBlob(string="foo"),
                    v=jobs_pb2.DataBlob(string="bar"),
                ),
            ]
        ),
    )
    self.assertEmpty(states[1].plugin_state.dat)

    self.db.UpdateHuntOutputPluginState(hunt_id, 1, Update)

    states = self.db.ReadHuntOutputPluginsStates(hunt_id)
    self.assertEqual(
        states[0].plugin_state,
        jobs_pb2.AttributedDict(
            dat=[
                jobs_pb2.KeyValue(
                    k=jobs_pb2.DataBlob(string="foo"),
                    v=jobs_pb2.DataBlob(string="bar"),
                ),
            ]
        ),
    )

    self.assertEqual(
        states[1].plugin_state,
        jobs_pb2.AttributedDict(
            dat=[
                jobs_pb2.KeyValue(
                    k=jobs_pb2.DataBlob(string="foo"),
                    v=jobs_pb2.DataBlob(string="bar"),
                ),
            ]
        ),
    )

  def testReadHuntLogEntriesReturnsEntryFromSingleHuntFlow(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(
        client_id="C.12345678901234aa", hunt_id=hunt_id
    )
    self.db.WriteFlowLogEntry(
        flows_pb2.FlowLogEntry(
            client_id=client_id,
            flow_id=flow_id,
            hunt_id=hunt_id,
            message="blah",
        )
    )

    hunt_log_entries = self.db.ReadHuntLogEntries(hunt_id, 0, 10)
    self.assertLen(hunt_log_entries, 1)
    self.assertIsInstance(hunt_log_entries[0], flows_pb2.FlowLogEntry)
    self.assertEqual(hunt_log_entries[0].hunt_id, hunt_id)
    self.assertEqual(hunt_log_entries[0].client_id, client_id)
    self.assertEqual(hunt_log_entries[0].flow_id, flow_id)
    self.assertEqual(hunt_log_entries[0].message, "blah")

  def _WriteNestedAndNonNestedLogEntries(self, hunt_id: str):
    client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
    # Top-level hunt-induced flows should have the hunt's ID.
    self.db.WriteFlowLogEntry(
        flows_pb2.FlowLogEntry(
            client_id=client_id,
            flow_id=flow_id,
            hunt_id=hunt_id,
            message="blah_a",
        )
    )
    self.db.WriteFlowLogEntry(
        flows_pb2.FlowLogEntry(
            client_id=client_id,
            flow_id=flow_id,
            hunt_id=hunt_id,
            message="blah_b",
        )
    )

    for i in range(10):
      _, nested_flow_id = self._SetupHuntClientAndFlow(
          client_id=client_id,
          parent_flow_id=flow_id,
          hunt_id=hunt_id,
          flow_id=flow.RandomFlowId(),
      )
      self.db.WriteFlowLogEntry(
          flows_pb2.FlowLogEntry(
              client_id=client_id,
              flow_id=nested_flow_id,
              hunt_id=hunt_id,
              message="blah_a_%d" % i,
          )
      )
      self.db.WriteFlowLogEntry(
          flows_pb2.FlowLogEntry(
              client_id=client_id,
              flow_id=nested_flow_id,
              hunt_id=hunt_id,
              message="blah_b_%d" % i,
          )
      )

  def testReadHuntLogEntriesIgnoresNestedFlows(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    self._WriteNestedAndNonNestedLogEntries(hunt_id)

    hunt_log_entries = self.db.ReadHuntLogEntries(hunt_id, 0, 10)
    self.assertLen(hunt_log_entries, 2)
    self.assertEqual(hunt_log_entries[0].message, "blah_a")
    self.assertEqual(hunt_log_entries[1].message, "blah_b")

  def testCountHuntLogEntriesIgnoresNestedFlows(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    self._WriteNestedAndNonNestedLogEntries(hunt_id)

    num_hunt_log_entries = self.db.CountHuntLogEntries(hunt_id)
    self.assertEqual(num_hunt_log_entries, 2)

  def _WriteHuntLogEntries(self, msg="blah") -> str:
    hunt_id = db_test_utils.InitializeHunt(self.db)

    for i in range(10):
      client_id, flow_id = self._SetupHuntClientAndFlow(
          client_id="C.12345678901234a%d" % i, hunt_id=hunt_id
      )
      self.db.WriteFlowLogEntry(
          flows_pb2.FlowLogEntry(
              client_id=client_id,
              flow_id=flow_id,
              hunt_id=hunt_id,
              message="%s%d" % (msg, i),
          )
      )

    return hunt_id

  def testReadHuntLogEntriesReturnsEntryFromMultipleHuntFlows(self):
    hunt_id = self._WriteHuntLogEntries()

    hunt_log_entries = self.db.ReadHuntLogEntries(hunt_id, 0, 100)
    self.assertLen(hunt_log_entries, 10)
    # Make sure messages are returned in timestamps-ascending order.
    for i, e in enumerate(hunt_log_entries):
      self.assertEqual(e.message, "blah%d" % i)

  def testReadHuntLogEntriesCorrectlyAppliesOffsetAndCountFilters(self):
    hunt_id = self._WriteHuntLogEntries()

    for i in range(10):
      hunt_log_entries = self.db.ReadHuntLogEntries(hunt_id, i, 1)
      self.assertLen(hunt_log_entries, 1)
      self.assertEqual(hunt_log_entries[0].message, "blah%d" % i)

  def testReadHuntLogEntriesCorrectlyAppliesWithSubstringFilter(self):
    hunt_id = self._WriteHuntLogEntries()

    hunt_log_entries = self.db.ReadHuntLogEntries(
        hunt_id, 0, 100, with_substring="foo"
    )
    self.assertEmpty(hunt_log_entries)

    hunt_log_entries = self.db.ReadHuntLogEntries(
        hunt_id, 0, 100, with_substring="blah"
    )
    self.assertLen(hunt_log_entries, 10)
    # Make sure messages are returned in timestamps-ascending order.
    for i, e in enumerate(hunt_log_entries):
      self.assertEqual(e.message, "blah%d" % i)

    hunt_log_entries = self.db.ReadHuntLogEntries(
        hunt_id, 0, 100, with_substring="blah1"
    )
    self.assertLen(hunt_log_entries, 1)
    self.assertEqual(hunt_log_entries[0].message, "blah1")

  def testReadHuntLogEntriesSubstringFilterIsCorrectlyEscaped(self):
    hunt_id = self._WriteHuntLogEntries("ABC%1")
    hunt_log_entries = self.db.ReadHuntLogEntries(
        hunt_id, 0, 100, with_substring="BC%1"
    )
    self.assertLen(hunt_log_entries, 10)
    hunt_log_entries = self.db.ReadHuntLogEntries(
        hunt_id, 0, 100, with_substring="B%1"
    )
    self.assertEmpty(hunt_log_entries)

  def testReadHuntLogEntriesCorrectlyAppliesCombinationOfFilters(self):
    hunt_id = self._WriteHuntLogEntries()

    hunt_log_entries = self.db.ReadHuntLogEntries(
        hunt_id, 0, 1, with_substring="blah"
    )
    self.assertLen(hunt_log_entries, 1)
    self.assertEqual(hunt_log_entries[0].message, "blah0")

  def testCountHuntLogEntriesReturnsCorrectHuntLogEntriesCount(self):
    hunt_id = self._WriteHuntLogEntries()

    num_entries = self.db.CountHuntLogEntries(hunt_id)
    self.assertEqual(num_entries, 10)

  def _WriteHuntResults(self, sample_results: List[flows_pb2.FlowResult]):
    self.db.WriteFlowResults(sample_results)

    # Update num_replies_sent for all flows referenced in sample_results:
    # in case the DB implementation relies on this data when
    # counting results.
    results_per_flow = collections.Counter()
    for r in sample_results:
      results_per_flow[(r.client_id, r.flow_id)] += 1

    for (client_id, flow_id), delta in results_per_flow.items():
      f_obj = self.db.ReadFlowObject(client_id, flow_id)
      f_obj.num_replies_sent += delta
      self.db.UpdateFlow(client_id, flow_id, flow_obj=f_obj)

  def _SampleSingleTypeHuntResults(
      self,
      client_id=None,
      flow_id=None,
      hunt_id=None,
      serial_number=None,
      count=10,
  ) -> List[flows_pb2.FlowResult]:
    self.assertIsNotNone(client_id)
    self.assertIsNotNone(flow_id)
    self.assertIsNotNone(hunt_id)

    res = []
    for i in range(count):
      payload_any = any_pb2.Any()
      payload_any.Pack(
          jobs_pb2.ClientSummary(
              client_id=client_id,
              system_manufacturer="manufacturer_%d" % i,
              serial_number=serial_number,
              install_date=int(
                  rdfvalue.RDFDatetime.FromSecondsSinceEpoch(10 + i)
              ),
          ),
      )
      res.append(
          flows_pb2.FlowResult(
              client_id=client_id,
              flow_id=flow_id,
              hunt_id=hunt_id,
              tag="tag_%d" % i,
              payload=payload_any,
          ),
      )
    return res

  def _SampleTwoTypeHuntResults(
      self,
      client_id=None,
      flow_id=None,
      hunt_id=None,
      serial_number=None,
      count_per_type=5,
      timestamp_start=10,
  ) -> List[flows_pb2.FlowResult]:
    self.assertIsNotNone(client_id)
    self.assertIsNotNone(flow_id)
    self.assertIsNotNone(hunt_id)

    res = []
    for i in range(count_per_type):
      paylad_any = any_pb2.Any()
      paylad_any.Pack(
          jobs_pb2.ClientSummary(
              client_id=client_id,
              system_manufacturer="manufacturer_%d" % i,
              serial_number=serial_number,
              install_date=int(
                  rdfvalue.RDFDatetime.FromSecondsSinceEpoch(
                      timestamp_start + i
                  )
              ),
          ),
      )
      res.append(
          flows_pb2.FlowResult(
              client_id=client_id,
              flow_id=flow_id,
              hunt_id=hunt_id,
              tag="tag_%d" % i,
              payload=paylad_any,
          )
      )
    for i in range(count_per_type):
      paylad_any = any_pb2.Any()
      paylad_any.Pack(
          jobs_pb2.ClientCrash(
              client_id=client_id,
              timestamp=int(
                  rdfvalue.RDFDatetime.FromSecondsSinceEpoch(
                      timestamp_start + i
                  )
              ),
          ),
      )
      res.append(
          flows_pb2.FlowResult(
              client_id=client_id,
              flow_id=flow_id,
              hunt_id=hunt_id,
              tag="tag_%d" % i,
              payload=paylad_any,
          )
      )
    return res

  def testReadHuntResultsReadsSingleResultOfSingleType(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
    sample_results = self._SampleSingleTypeHuntResults(
        client_id=client_id, flow_id=flow_id, hunt_id=hunt_id, count=1
    )
    self._WriteHuntResults(sample_results)

    results = self.db.ReadHuntResults(hunt_id, 0, 10)
    self.assertLen(results, 1)
    self.assertEqual(results[0].hunt_id, hunt_id)
    self.assertEqual(results[0].payload, sample_results[0].payload)

  def testReadHuntResultsReadsMultipleResultOfSingleType(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
    sample_results = self._SampleSingleTypeHuntResults(
        client_id=client_id, flow_id=flow_id, hunt_id=hunt_id, count=10
    )
    self._WriteHuntResults(sample_results)

    results = self.db.ReadHuntResults(hunt_id, 0, 1000)
    self.assertLen(results, 10)
    for i in range(10):
      self.assertEqual(results[i].hunt_id, hunt_id)
      self.assertEqual(results[i].payload, sample_results[i].payload)

  def testReadHuntResultsReadsMultipleResultOfMultipleTypes(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id_1, flow_id_1 = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
    sample_results_1 = self._SampleTwoTypeHuntResults(
        client_id=client_id_1, flow_id=flow_id_1, hunt_id=hunt_id
    )
    self._WriteHuntResults(sample_results_1)

    client_id_2, flow_id_2 = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
    sample_results_2 = self._SampleTwoTypeHuntResults(
        client_id=client_id_2, flow_id=flow_id_2, hunt_id=hunt_id
    )
    self._WriteHuntResults(sample_results_2)

    sample_results = sample_results_1 + sample_results_2
    results = self.db.ReadHuntResults(hunt_id, 0, 1000)
    self.assertLen(results, len(sample_results))
    self.assertListEqual(
        [i.payload for i in results], [i.payload for i in sample_results]
    )

  def testReadHuntResultsCorrectlyAppliedOffsetAndCountFilters(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    sample_results = []
    for _ in range(10):
      client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
      results = self._SampleSingleTypeHuntResults(
          client_id=client_id, flow_id=flow_id, hunt_id=hunt_id, count=1
      )
      sample_results.extend(results)
      self._WriteHuntResults(results)

    for l in range(1, 11):
      for i in range(10):
        results = self.db.ReadHuntResults(hunt_id, i, l)
        expected = sample_results[i : i + l]

        result_payloads = [x.payload for x in results]
        expected_payloads = [x.payload for x in expected]
        self.assertEqual(
            result_payloads,
            expected_payloads,
            "Results differ from expected (from %d, size %d): %s vs %s"
            % (i, l, result_payloads, expected_payloads),
        )

  def testReadHuntResultsCorrectlyAppliesWithTagFilter(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
    sample_results = self._SampleSingleTypeHuntResults(
        client_id=client_id, flow_id=flow_id, hunt_id=hunt_id
    )
    self._WriteHuntResults(sample_results)

    results = self.db.ReadHuntResults(hunt_id, 0, 100, with_tag="blah")
    self.assertFalse(results)

    results = self.db.ReadHuntResults(hunt_id, 0, 100, with_tag="tag")
    self.assertFalse(results)

    results = self.db.ReadHuntResults(hunt_id, 0, 100, with_tag="tag_1")
    self.assertEqual(
        [i.payload for i in results],
        [i.payload for i in sample_results if i.tag == "tag_1"],
    )

  def testReadHuntResultsCorrectlyAppliesWithTypeFilter(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    sample_results = []
    for _ in range(10):
      client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
      results = self._SampleTwoTypeHuntResults(
          client_id=client_id,
          flow_id=flow_id,
          hunt_id=hunt_id,
          count_per_type=1,
      )
      sample_results.extend(results)
      self._WriteHuntResults(results)

    results = self.db.ReadHuntResults(
        hunt_id, 0, 100, with_type=jobs_pb2.ClientInformation.__name__
    )
    self.assertFalse(results)

    results = self.db.ReadHuntResults(
        hunt_id, 0, 100, with_type=jobs_pb2.ClientSummary.__name__
    )
    self.assertCountEqual(
        [i.payload for i in results],
        [
            i.payload
            for i in sample_results
            if i.payload.Is(jobs_pb2.ClientSummary.DESCRIPTOR)
        ],
    )

  def testReadHuntResultsCorrectlyAppliesWithSubstringFilter(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
    sample_results = self._SampleSingleTypeHuntResults(
        client_id=client_id, flow_id=flow_id, hunt_id=hunt_id
    )
    self._WriteHuntResults(sample_results)

    results = self.db.ReadHuntResults(hunt_id, 0, 100, with_substring="blah")
    self.assertEmpty(results)

    results = self.db.ReadHuntResults(
        hunt_id, 0, 100, with_substring="manufacturer"
    )
    self.assertEqual(
        [i.payload for i in results],
        [i.payload for i in sample_results],
    )

    results = self.db.ReadHuntResults(
        hunt_id, 0, 100, with_substring="manufacturer_1"
    )
    self.assertEqual([i.payload for i in results], [sample_results[1].payload])

  def testReadHuntResultsSubstringFilterIsCorrectlyEscaped(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)

    sample_results = self._SampleSingleTypeHuntResults(
        client_id=client_id,
        flow_id=flow_id,
        hunt_id=hunt_id,
        serial_number="ABC%123",
    )
    self._WriteHuntResults(sample_results)

    results = self.db.ReadHuntResults(hunt_id, 0, 100, with_substring="ABC%123")
    self.assertLen(results, 10)

    results = self.db.ReadHuntResults(hunt_id, 0, 100, with_substring="AB%23")
    self.assertEmpty(results)

  def testReadHuntResultsCorrectlyAppliesVariousCombinationsOfFilters(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    sample_results = []
    for _ in range(10):
      client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
      results = self._SampleTwoTypeHuntResults(
          client_id=client_id,
          flow_id=flow_id,
          hunt_id=hunt_id,
          count_per_type=5,
      )
      sample_results.extend(results)
      self._WriteHuntResults(results)

    # TODO: Clean up this test.
    tags = {
        None: list(sample_results),
        "tag_1": [s for s in sample_results if s.tag == "tag_1"],
    }
    substrings = {
        None: list(sample_results),
    }
    manufacturer = []
    manufacturer_1 = []
    for s in sample_results:
      if s.payload.Is(jobs_pb2.ClientSummary.DESCRIPTOR):
        payload = jobs_pb2.ClientSummary()
      elif s.payload.Is(jobs_pb2.ClientCrash.DESCRIPTOR):
        payload = jobs_pb2.ClientCrash()
      else:
        continue
      s.payload.Unpack(payload)
      if "manufacturer" in getattr(payload, "system_manufacturer", ""):
        manufacturer.append(s)
      if "manufacturer_1" in getattr(payload, "system_manufacturer", ""):
        manufacturer_1.append(s)

    substrings["manufacturer"] = manufacturer
    substrings["manufacturer_1"] = manufacturer_1

    types = {
        None: list(sample_results),
        jobs_pb2.ClientSummary.__name__: [
            s
            for s in sample_results
            if s.payload.Is(jobs_pb2.ClientSummary.DESCRIPTOR)
        ],
    }

    for tag_value, tag_expected in tags.items():
      for substring_value, substring_expected in substrings.items():
        for type_value, type_expected in types.items():
          expected = [
              e
              for e in tag_expected
              if e in substring_expected and e in type_expected
          ]
          results = self.db.ReadHuntResults(
              hunt_id,
              0,
              100,
              with_tag=tag_value,
              with_type=type_value,
              with_substring=substring_value,
          )

          self.assertCountEqual(
              [i.payload for i in expected],
              [i.payload for i in results],
              "Result items do not match for "
              "(tag=%s, type=%s, substring=%s): %s vs %s"
              % (tag_value, type_value, substring_value, expected, results),
          )

  def testReadHuntResultsReturnsPayloadWithMissingTypeAsSpecialValue(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
    sample_results = self._SampleSingleTypeHuntResults(
        client_id=client_id, flow_id=flow_id, hunt_id=hunt_id
    )
    self._WriteHuntResults(sample_results)

    type_name = rdf_client.ClientSummary.__name__
    cls = rdfvalue.RDFValue.classes.pop(type_name)

    results = self.db.ReadHuntResults(hunt_id, 0, 100)
    rdfvalue.RDFValue.classes[type_name] = cls

    self.assertLen(sample_results, len(results))
    for r in results:
      self.assertTrue(
          r.payload.Is(objects_pb2.SerializedValueOfUnrecognizedType.DESCRIPTOR)
      )
      payload = objects_pb2.SerializedValueOfUnrecognizedType()
      r.payload.Unpack(payload)
      self.assertEqual(payload.type_name, type_name)

  def testReadHuntResultsIgnoresChildFlowsResults(self):
    client_id = db_test_utils.InitializeClient(self.db)
    hunt_id = db_test_utils.InitializeHunt(self.db)

    hunt_flow_id = db_test_utils.InitializeFlow(
        self.db,
        client_id=client_id,
        flow_id=hunt_id,
        parent_hunt_id=hunt_id,
    )

    child_flow_id = db_test_utils.InitializeFlow(
        self.db,
        client_id=client_id,
        parent_flow_id=hunt_flow_id,
        parent_hunt_id=hunt_id,
    )

    hunt_flow_result = flows_pb2.FlowResult()
    hunt_flow_result.client_id = client_id
    hunt_flow_result.hunt_id = hunt_id
    hunt_flow_result.flow_id = hunt_flow_id
    hunt_flow_result.payload.Pack(jobs_pb2.Uname(fqdn="hunt.example.com"))
    self.db.WriteFlowResults([hunt_flow_result])

    child_flow_result = flows_pb2.FlowResult()
    child_flow_result.client_id = client_id
    child_flow_result.hunt_id = hunt_id
    child_flow_result.flow_id = child_flow_id
    child_flow_result.payload.Pack(jobs_pb2.Uname(fqdn="child.example.com"))
    self.db.WriteFlowResults([child_flow_result])

    results = self.db.ReadHuntResults(hunt_id, offset=0, count=1024)

    self.assertLen(results, 1)

    result = jobs_pb2.Uname()
    results[0].payload.Unpack(result)
    self.assertEqual(result.fqdn, "hunt.example.com")

  def testCountHuntResultsReturnsCorrectResultsCount(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
    sample_results = self._SampleSingleTypeHuntResults(
        client_id=client_id, flow_id=flow_id, hunt_id=hunt_id
    )
    self._WriteHuntResults(sample_results)

    num_results = self.db.CountHuntResults(hunt_id)
    self.assertLen(sample_results, num_results)

  def testCountHuntResultsCorrectlyAppliesWithTagFilter(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
    sample_results = self._SampleSingleTypeHuntResults(
        client_id=client_id, flow_id=flow_id, hunt_id=hunt_id
    )
    self._WriteHuntResults(sample_results)

    num_results = self.db.CountHuntResults(hunt_id, with_tag="blah")
    self.assertEqual(num_results, 0)

    num_results = self.db.CountHuntResults(hunt_id, with_tag="tag_1")
    self.assertEqual(num_results, 1)

  def testCountHuntResultsCorrectlyAppliesWithTypeFilter(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    sample_results = []
    for _ in range(10):
      client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
      results = self._SampleTwoTypeHuntResults(
          client_id=client_id,
          flow_id=flow_id,
          hunt_id=hunt_id,
          count_per_type=1,
      )
      sample_results.extend(results)
      self._WriteHuntResults(results)

    num_results = self.db.CountHuntResults(
        hunt_id, with_type=rdf_client.ClientInformation.__name__
    )
    self.assertEqual(num_results, 0)

    num_results = self.db.CountHuntResults(
        hunt_id, with_type=rdf_client.ClientSummary.__name__
    )
    self.assertEqual(num_results, 10)

    num_results = self.db.CountHuntResults(
        hunt_id, with_type=rdf_client.ClientCrash.__name__
    )
    self.assertEqual(num_results, 10)

  def testCountHuntResultsCorrectlyAppliesWithTagAndWithTypeFilters(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    sample_results = []
    for _ in range(10):
      client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
      results = self._SampleTwoTypeHuntResults(
          client_id=client_id,
          flow_id=flow_id,
          hunt_id=hunt_id,
          count_per_type=5,
      )
      sample_results.extend(results)
      self._WriteHuntResults(results)

    num_results = self.db.CountHuntResults(
        hunt_id, with_tag="tag_1", with_type=rdf_client.ClientSummary.__name__
    )
    self.assertEqual(num_results, 10)

  def testCountHuntResultsCorrectlyAppliesWithTimestampFilter(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    for _ in range(10):
      client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
      sample_results = self._SampleSingleTypeHuntResults(
          client_id=client_id, flow_id=flow_id, hunt_id=hunt_id, count=10
      )
      self._WriteHuntResults(sample_results[:5])
      self._WriteHuntResults(sample_results[5:])

    hunt_results = self.db.ReadHuntResults(hunt_id, 0, 10)

    for hr in hunt_results:
      self.assertEqual(
          [hr],
          self.db.ReadHuntResults(
              hunt_id,
              0,
              10,
              with_timestamp=rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
                  hr.timestamp
              ),
          ),
      )

  def testCountHuntResultsByTypeGroupsResultsCorrectly(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
    results = self._SampleTwoTypeHuntResults(
        client_id=client_id, flow_id=flow_id, hunt_id=hunt_id, count_per_type=5
    )
    self._WriteHuntResults(results)

    counts = self.db.CountHuntResultsByType(hunt_id)
    for key in counts:
      self.assertIsInstance(key, str)

    self.assertEqual(
        counts,
        {
            rdf_client.ClientSummary.__name__: 5,
            rdf_client.ClientCrash.__name__: 5,
        },
    )

  def testReadHuntFlowsReturnsEmptyListWhenNoFlows(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    self.assertEmpty(self.db.ReadHuntFlows(hunt_id, 0, 10))

  def testReadHuntFlowsReturnsAllHuntFlowsWhenNoFilterCondition(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    _, flow_id_1 = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
    _, flow_id_2 = self._SetupHuntClientAndFlow(hunt_id=hunt_id)

    flows = self.db.ReadHuntFlows(hunt_id, 0, 10)
    self.assertCountEqual([f.flow_id for f in flows], [flow_id_1, flow_id_2])

  def _BuildFilterConditionExpectations(self, hunt_id):
    _, running_flow_id = self._SetupHuntClientAndFlow(
        flow_state=rdf_flow_objects.Flow.FlowState.RUNNING, hunt_id=hunt_id
    )
    _, succeeded_flow_id = self._SetupHuntClientAndFlow(
        flow_state=rdf_flow_objects.Flow.FlowState.FINISHED, hunt_id=hunt_id
    )
    _, failed_flow_id = self._SetupHuntClientAndFlow(
        flow_state=rdf_flow_objects.Flow.FlowState.ERROR, hunt_id=hunt_id
    )
    _, crashed_flow_id = self._SetupHuntClientAndFlow(
        flow_state=rdf_flow_objects.Flow.FlowState.CRASHED, hunt_id=hunt_id
    )
    client_id, flow_with_results_id = self._SetupHuntClientAndFlow(
        hunt_id=hunt_id
    )
    sample_results = self._SampleSingleTypeHuntResults(
        client_id=client_id, flow_id=flow_with_results_id, hunt_id=hunt_id
    )
    self._WriteHuntResults(sample_results)

    return {
        db.HuntFlowsCondition.UNSET: [
            running_flow_id,
            succeeded_flow_id,
            failed_flow_id,
            crashed_flow_id,
            flow_with_results_id,
        ],
        db.HuntFlowsCondition.FAILED_FLOWS_ONLY: [failed_flow_id],
        db.HuntFlowsCondition.SUCCEEDED_FLOWS_ONLY: [succeeded_flow_id],
        db.HuntFlowsCondition.COMPLETED_FLOWS_ONLY: [
            failed_flow_id,
            succeeded_flow_id,
        ],
        db.HuntFlowsCondition.FLOWS_IN_PROGRESS_ONLY: [running_flow_id],
        db.HuntFlowsCondition.CRASHED_FLOWS_ONLY: [crashed_flow_id],
    }

  def testReadHuntFlowsAppliesFilterConditionCorrectly(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    expectations = self._BuildFilterConditionExpectations(hunt_id)
    for filter_condition, expected in expectations.items():
      results = self.db.ReadHuntFlows(
          hunt_id, 0, 10, filter_condition=filter_condition
      )
      results_ids = [r.flow_id for r in results]
      self.assertCountEqual(
          results_ids,
          expected,
          "Result items do not match for (filter_condition=%s): %s vs %s"
          % (filter_condition, expected, results_ids),
      )

  def testReadHuntFlowsCorrectlyAppliesOffsetAndCountFilters(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    expectations = self._BuildFilterConditionExpectations(hunt_id)
    for filter_condition, _ in expectations.items():
      full_results = self.db.ReadHuntFlows(
          hunt_id, 0, 1024, filter_condition=filter_condition
      )
      full_results_ids = [r.flow_id for r in full_results]
      for index in range(0, 2):
        for count in range(1, 3):
          results = self.db.ReadHuntFlows(
              hunt_id, index, count, filter_condition=filter_condition
          )
          results_ids = [r.flow_id for r in results]
          expected_ids = full_results_ids[index : index + count]
          self.assertCountEqual(
              results_ids,
              expected_ids,
              "Result items do not match for "
              "(filter_condition=%s, index=%d, count=%d): %s vs %s"
              % (filter_condition, index, count, expected_ids, results_ids),
          )

  def testReadHuntFlowsIgnoresSubflows(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(
        hunt_id=hunt_id, flow_state=rdf_flow_objects.Flow.FlowState.RUNNING
    )

    # Whatever state the subflow is in, it should be ignored.
    self._SetupHuntClientAndFlow(
        client_id=client_id,
        hunt_id=hunt_id,
        flow_id=flow.RandomFlowId(),
        parent_flow_id=flow_id,
        flow_state=rdf_flow_objects.Flow.FlowState.ERROR,
    )
    self._SetupHuntClientAndFlow(
        client_id=client_id,
        hunt_id=hunt_id,
        flow_id=flow.RandomFlowId(),
        parent_flow_id=flow_id,
        flow_state=rdf_flow_objects.Flow.FlowState.FINISHED,
    )
    self._SetupHuntClientAndFlow(
        client_id=client_id,
        hunt_id=hunt_id,
        flow_id=flow.RandomFlowId(),
        parent_flow_id=flow_id,
        flow_state=rdf_flow_objects.Flow.FlowState.RUNNING,
    )

    for state, expected_results in [
        (db.HuntFlowsCondition.COMPLETED_FLOWS_ONLY, 0),
        (db.HuntFlowsCondition.SUCCEEDED_FLOWS_ONLY, 0),
        (db.HuntFlowsCondition.FLOWS_IN_PROGRESS_ONLY, 1),
    ]:
      results = self.db.ReadHuntFlows(hunt_id, 0, 10, filter_condition=state)
      self.assertLen(results, expected_results)

  def testCountHuntFlowsIgnoresSubflows(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(
        hunt_id=hunt_id, flow_state=rdf_flow_objects.Flow.FlowState.RUNNING
    )

    # Whatever state the subflow is in, it should be ignored.
    self._SetupHuntClientAndFlow(
        client_id=client_id,
        hunt_id=hunt_id,
        flow_id=flow.RandomFlowId(),
        parent_flow_id=flow_id,
        flow_state=rdf_flow_objects.Flow.FlowState.ERROR,
    )
    self._SetupHuntClientAndFlow(
        client_id=client_id,
        hunt_id=hunt_id,
        flow_id=flow.RandomFlowId(),
        parent_flow_id=flow_id,
        flow_state=rdf_flow_objects.Flow.FlowState.FINISHED,
    )
    self._SetupHuntClientAndFlow(
        client_id=client_id,
        hunt_id=hunt_id,
        flow_id=flow.RandomFlowId(),
        parent_flow_id=flow_id,
        flow_state=rdf_flow_objects.Flow.FlowState.RUNNING,
    )

    self.assertEqual(self.db.CountHuntFlows(hunt_id), 1)

  def testCountHuntFlowsReturnsEmptyListWhenNoFlows(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    self.assertEqual(self.db.CountHuntFlows(hunt_id), 0)

  def testCountHuntFlowsReturnsAllHuntFlowsWhenNoFilterCondition(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    self._SetupHuntClientAndFlow(hunt_id=hunt_id)
    self._SetupHuntClientAndFlow(hunt_id=hunt_id)

    self.assertEqual(self.db.CountHuntFlows(hunt_id), 2)

  def testCountHuntFlowsAppliesFilterConditionCorrectly(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    expectations = self._BuildFilterConditionExpectations(hunt_id)
    for filter_condition, expected in expectations.items():
      result = self.db.CountHuntFlows(
          hunt_id, filter_condition=filter_condition
      )
      self.assertLen(
          expected,
          result,
          "Result count does not match for (filter_condition=%s): %d vs %d"
          % (filter_condition, len(expected), result),
      )

  def testReadHuntFlowErrors(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id_1 = db_test_utils.InitializeClient(self.db)
    client_id_2 = db_test_utils.InitializeClient(self.db)

    flow_id_1 = db_test_utils.InitializeFlow(
        self.db,
        client_id=client_id_1,
        flow_id=hunt_id,
        parent_hunt_id=hunt_id,
    )
    flow_id_2 = db_test_utils.InitializeFlow(
        self.db,
        client_id=client_id_2,
        flow_id=hunt_id,
        parent_hunt_id=hunt_id,
    )

    pre_update_1_time = self.db.Now()

    flow_obj_1 = self.db.ReadFlowObject(client_id_1, flow_id_1)
    flow_obj_1.flow_state = flows_pb2.Flow.FlowState.ERROR
    flow_obj_1.error_message = "ERROR_1"
    flow_obj_1.backtrace = "File 'foo.py', line 1, in 'foo'"
    self.db.UpdateFlow(client_id_1, flow_id_1, flow_obj_1)

    pre_update_2_time = self.db.Now()

    flow_obj_2 = self.db.ReadFlowObject(client_id_2, flow_id_2)
    flow_obj_2.flow_state = flows_pb2.Flow.FlowState.ERROR
    flow_obj_2.error_message = "ERROR_2"
    self.db.UpdateFlow(client_id_2, flow_id_2, flow_obj_2)

    results = self.db.ReadHuntFlowErrors(hunt_id, offset=0, count=1024)
    self.assertLen(results, 2)

    self.assertEqual(results[client_id_1].message, "ERROR_1")
    self.assertGreater(results[client_id_1].time, pre_update_1_time)
    self.assertIsNotNone(results[client_id_1].backtrace)

    self.assertEqual(results[client_id_2].message, "ERROR_2")
    self.assertGreater(results[client_id_2].time, pre_update_2_time)
    self.assertIsNone(results[client_id_2].backtrace)

  def testReadHuntFlowErrorsIgnoreSubflows(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id = db_test_utils.InitializeClient(self.db)

    flow_id = db_test_utils.InitializeFlow(
        self.db,
        client_id=client_id,
        flow_id=hunt_id,
        parent_hunt_id=hunt_id,
    )
    child_flow_id = db_test_utils.InitializeFlow(
        self.db,
        client_id=client_id,
        parent_flow_id=flow_id,
        parent_hunt_id=hunt_id,
    )

    flow_obj = self.db.ReadFlowObject(client_id, flow_id)
    flow_obj.flow_state = flows_pb2.Flow.FlowState.ERROR
    flow_obj.error_message = "ERROR"
    self.db.UpdateFlow(client_id, flow_id, flow_obj)

    child_flow_obj = self.db.ReadFlowObject(client_id, child_flow_id)
    child_flow_obj.flow_state = flows_pb2.Flow.FlowState.ERROR
    child_flow_obj.error_message = "CHILD_ERROR"
    self.db.UpdateFlow(client_id, child_flow_id, child_flow_obj)

    results = self.db.ReadHuntFlowErrors(hunt_id, offset=0, count=1024)
    self.assertLen(results, 1)
    self.assertEqual(results[client_id].message, "ERROR")

  def testReadHuntFlowErrorsOffsetAndCount(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id_1 = db_test_utils.InitializeClient(self.db)
    client_id_2 = db_test_utils.InitializeClient(self.db)
    client_id_3 = db_test_utils.InitializeClient(self.db)

    flow_id_1 = db_test_utils.InitializeFlow(
        self.db,
        client_id=client_id_1,
        flow_id=hunt_id,
        parent_hunt_id=hunt_id,
    )
    flow_id_2 = db_test_utils.InitializeFlow(
        self.db,
        client_id=client_id_2,
        flow_id=hunt_id,
        parent_hunt_id=hunt_id,
    )
    flow_id_3 = db_test_utils.InitializeFlow(
        self.db,
        client_id=client_id_3,
        flow_id=hunt_id,
        parent_hunt_id=hunt_id,
    )

    flow_obj_1 = self.db.ReadFlowObject(client_id_1, flow_id_1)
    flow_obj_1.flow_state = flows_pb2.Flow.FlowState.ERROR
    flow_obj_1.error_message = "ERROR_1"
    self.db.UpdateFlow(client_id_1, flow_id_1, flow_obj_1)

    flow_obj_2 = self.db.ReadFlowObject(client_id_2, flow_id_2)
    flow_obj_2.flow_state = flows_pb2.Flow.FlowState.ERROR
    flow_obj_2.error_message = "ERROR_2"
    self.db.UpdateFlow(client_id_2, flow_id_2, flow_obj_2)

    flow_obj_3 = self.db.ReadFlowObject(client_id_3, flow_id_3)
    flow_obj_3.flow_state = flows_pb2.Flow.FlowState.ERROR
    flow_obj_3.error_message = "ERROR_3"
    self.db.UpdateFlow(client_id_3, flow_id_3, flow_obj_3)

    results = self.db.ReadHuntFlowErrors(hunt_id, offset=1, count=1)
    self.assertLen(results, 1)
    self.assertEqual(results[client_id_2].message, "ERROR_2")

  def testReadHuntCountersForNewHunt(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)
    hunt_counters = self.db.ReadHuntCounters(hunt_id)
    self.assertEqual(hunt_counters.num_clients, 0)
    self.assertEqual(hunt_counters.num_successful_clients, 0)
    self.assertEqual(hunt_counters.num_failed_clients, 0)
    self.assertEqual(hunt_counters.num_clients_with_results, 0)
    self.assertEqual(hunt_counters.num_crashed_clients, 0)
    self.assertEqual(hunt_counters.num_running_clients, 0)
    self.assertEqual(hunt_counters.num_results, 0)
    self.assertEqual(hunt_counters.total_cpu_seconds, 0)
    self.assertEqual(hunt_counters.total_network_bytes_sent, 0)

  def testReadHuntsCountersForEmptyList(self):
    hunts_counters = self.db.ReadHuntsCounters([])
    self.assertEmpty(hunts_counters)

  def testReadHuntsCountersForSeveralHunts(self):
    hunt_id_1 = db_test_utils.InitializeHunt(self.db)
    hunt_id_2 = db_test_utils.InitializeHunt(self.db)

    hunts_counters = self.db.ReadHuntsCounters([hunt_id_1, hunt_id_2])
    self.assertLen(hunts_counters, 2)
    self.assertIsInstance(hunts_counters[hunt_id_1], db.HuntCounters)
    self.assertIsInstance(hunts_counters[hunt_id_2], db.HuntCounters)

  def testReadHuntsCountersForASubsetOfCreatedHunts(self):
    db_test_utils.InitializeUser(self.db, "user")

    hunt_id_1 = db_test_utils.InitializeHunt(self.db, creator="user")
    _ = db_test_utils.InitializeHunt(self.db, creator="user")
    hunt_id_3 = db_test_utils.InitializeHunt(self.db, creator="user")

    hunts_counters = self.db.ReadHuntsCounters([hunt_id_1, hunt_id_3])

    self.assertLen(hunts_counters, 2)
    self.assertIsInstance(hunts_counters[hunt_id_1], db.HuntCounters)
    self.assertIsInstance(hunts_counters[hunt_id_3], db.HuntCounters)

  def testReadHuntsCountersReturnsSameResultAsReadHuntCounters(self):
    db_test_utils.InitializeUser(self.db, "user")

    hunt_id_1 = db_test_utils.InitializeHunt(self.db, creator="user")
    self._BuildFilterConditionExpectations(hunt_id_1)

    hunt_id_2 = db_test_utils.InitializeHunt(self.db, creator="user")
    self._BuildFilterConditionExpectations(hunt_id_2)

    hunts_counters = self.db.ReadHuntsCounters([hunt_id_1, hunt_id_2])
    self.assertLen(hunts_counters, 2)
    self.assertEqual(
        hunts_counters[hunt_id_1],
        self.db.ReadHuntCounters(hunt_id_1),
    )
    self.assertEqual(
        hunts_counters[hunt_id_2],
        self.db.ReadHuntCounters(hunt_id_2),
    )

  def testReadHuntCountersCorrectlyAggregatesResultsAmongDifferentFlows(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    expectations = self._BuildFilterConditionExpectations(hunt_id)

    hunt_counters = self.db.ReadHuntCounters(hunt_id)
    self.assertLen(
        expectations[db.HuntFlowsCondition.UNSET], hunt_counters.num_clients
    )
    self.assertLen(
        expectations[db.HuntFlowsCondition.SUCCEEDED_FLOWS_ONLY],
        hunt_counters.num_successful_clients,
    )
    self.assertLen(
        expectations[db.HuntFlowsCondition.FAILED_FLOWS_ONLY],
        hunt_counters.num_failed_clients,
    )
    self.assertLen(
        expectations[db.HuntFlowsCondition.FLOWS_IN_PROGRESS_ONLY],
        hunt_counters.num_running_clients,
    )

    # _BuildFilterConditionExpectations writes 10 sample results for one client.
    self.assertEqual(hunt_counters.num_clients_with_results, 1)
    self.assertLen(
        expectations[db.HuntFlowsCondition.CRASHED_FLOWS_ONLY],
        hunt_counters.num_crashed_clients,
    )

    # _BuildFilterConditionExpectations writes 10 sample results.
    self.assertEqual(hunt_counters.num_results, 10)

    self.assertEqual(hunt_counters.total_cpu_seconds, 0)
    self.assertEqual(hunt_counters.total_network_bytes_sent, 0)

    # Check that after adding a flow with resource metrics, total counters
    # get updated.
    client_id = db_test_utils.InitializeClient(self.db)
    db_test_utils.InitializeFlow(
        self.db,
        client_id,
        flow_id=hunt_id,
        parent_hunt_id=hunt_id,
        cpu_time_used=rdf_client_stats.CpuSeconds(
            user_cpu_time=4.5, system_cpu_time=10
        ),
        network_bytes_sent=42,
    )

    hunt_counters = self.db.ReadHuntCounters(hunt_id)
    self.assertAlmostEqual(hunt_counters.total_cpu_seconds, 14.5)
    self.assertEqual(hunt_counters.total_network_bytes_sent, 42)

  def testReadHuntClientResourcesStatsIgnoresSubflows(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id = db_test_utils.InitializeClient(self.db)
    flow_id = db_test_utils.InitializeFlow(
        self.db,
        client_id,
        flow_id=hunt_id,
        flow_state=rdf_flow_objects.Flow.FlowState.FINISHED,
        parent_hunt_id=hunt_id,
        cpu_time_used=rdf_client_stats.CpuSeconds(
            user_cpu_time=100, system_cpu_time=200
        ),
        network_bytes_sent=300,
    )

    # Create a subflow that used some resources too. This resource usage is
    # already accounted for in the parent flow so the overall hunt resource
    # usage should ignore those numbers.
    sub_flow = rdf_flow_objects.Flow(
        client_id=client_id,
        flow_id="12345678",
        parent_flow_id=flow_id,
        parent_hunt_id=hunt_id,
        cpu_time_used=rdf_client_stats.CpuSeconds(
            user_cpu_time=10, system_cpu_time=20
        ),
        network_bytes_sent=30,
    )
    sub_flow = mig_flow_objects.ToProtoFlow(sub_flow)
    self.db.WriteFlowObject(sub_flow)

    usage_stats = self.db.ReadHuntClientResourcesStats(hunt_id)
    network_bins = usage_stats.network_bytes_sent_stats.histogram.bins
    user_cpu_bins = usage_stats.user_cpu_stats.histogram.bins
    system_cpu_bins = usage_stats.system_cpu_stats.histogram.bins

    self.assertEqual(usage_stats.network_bytes_sent_stats.sum, 300)
    self.assertEqual(sum([b.num for b in network_bins]), 1)
    self.assertEqual(usage_stats.system_cpu_stats.sum, 200)
    self.assertEqual(sum([b.num for b in system_cpu_bins]), 1)
    self.assertEqual(usage_stats.user_cpu_stats.sum, 100)
    self.assertEqual(sum([b.num for b in user_cpu_bins]), 1)
    self.assertLen(usage_stats.worst_performers, 1)

  def testReadHuntClientResourcesStatsCorrectlyAggregatesData(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    flow_data = []

    for i in range(10):
      user_cpu_time = 4.5 + i
      system_cpu_time = 10 + i * 2
      network_bytes_sent = 42 + i * 3

      client_id = db_test_utils.InitializeClient(self.db)
      flow_id = db_test_utils.InitializeFlow(
          self.db,
          client_id,
          flow_id=hunt_id,
          flow_state=rdf_flow_objects.Flow.FlowState.FINISHED,
          parent_hunt_id=hunt_id,
          cpu_time_used=rdf_client_stats.CpuSeconds(
              user_cpu_time=user_cpu_time, system_cpu_time=system_cpu_time
          ),
          network_bytes_sent=network_bytes_sent,
      )

      flow_data.append((
          client_id,
          flow_id,
          (user_cpu_time, system_cpu_time, network_bytes_sent),
      ))

    usage_stats = self.db.ReadHuntClientResourcesStats(hunt_id)

    expected_cpu_bins = model_hunts.CPU_STATS_BINS
    expected_network_bins = model_hunts.NETWORK_STATS_BINS
    expected_user_cpu_histogram = jobs_pb2.StatsHistogram(
        bins=[
            jobs_pb2.StatsHistogramBin(num=num, range_max_value=max_range)
            for num, max_range in zip(
                12 * [0] + [1, 1, 1, 1, 1, 1, 4, 0], expected_cpu_bins
            )
        ]
    )

    expected_system_cpu_histogram = jobs_pb2.StatsHistogram(
        bins=[
            jobs_pb2.StatsHistogramBin(num=num, range_max_value=max_range)
            for num, max_range in zip(18 * [0] + [3, 7], expected_cpu_bins)
        ]
    )
    expected_network_histogram = jobs_pb2.StatsHistogram(
        bins=[
            jobs_pb2.StatsHistogramBin(num=num, range_max_value=max_range)
            for num, max_range in zip(
                [0, 0, 8, 2] + 14 * [0], expected_network_bins
            )
        ]
    )

    self.assertEqual(usage_stats.user_cpu_stats.num, 10)
    self.assertAlmostEqual(usage_stats.user_cpu_stats.sum, 90)
    self.assertAlmostEqual(usage_stats.user_cpu_stats.stddev, 2.87228, 5)
    self.assertLen(
        usage_stats.user_cpu_stats.histogram.bins,
        len(expected_user_cpu_histogram.bins),
    )
    for bin_index, (b, exp_b) in enumerate(
        zip(
            usage_stats.user_cpu_stats.histogram.bins,
            expected_user_cpu_histogram.bins,
        )
    ):
      self.assertAlmostEqual(
          b.range_max_value, exp_b.range_max_value, msg=f"bin index {bin_index}"
      )
      self.assertEqual(b.num, exp_b.num, msg=f"bin index {bin_index}")

    self.assertEqual(usage_stats.system_cpu_stats.num, 10)
    self.assertAlmostEqual(usage_stats.system_cpu_stats.sum, 190)
    self.assertAlmostEqual(usage_stats.system_cpu_stats.stddev, 5.74456, 5)
    self.assertLen(
        usage_stats.system_cpu_stats.histogram.bins,
        len(expected_system_cpu_histogram.bins),
    )
    for bin_index, (b, exp_b) in enumerate(
        zip(
            usage_stats.system_cpu_stats.histogram.bins,
            expected_system_cpu_histogram.bins,
        )
    ):
      self.assertAlmostEqual(
          b.range_max_value, exp_b.range_max_value, msg=f"bin index {bin_index}"
      )
      self.assertEqual(b.num, exp_b.num, msg=f"bin index {bin_index}")

    self.assertEqual(usage_stats.network_bytes_sent_stats.num, 10)
    self.assertAlmostEqual(usage_stats.network_bytes_sent_stats.sum, 555)
    self.assertAlmostEqual(
        usage_stats.network_bytes_sent_stats.stddev, 8.6168, 4
    )
    self.assertLen(
        usage_stats.network_bytes_sent_stats.histogram.bins,
        len(expected_network_histogram.bins),
    )
    for bin_index, (b, model_b) in enumerate(
        zip(
            usage_stats.network_bytes_sent_stats.histogram.bins,
            expected_network_histogram.bins,
        )
    ):
      self.assertAlmostEqual(
          b.range_max_value,
          model_b.range_max_value,
          msg=f"bin index {bin_index}",
      )
      self.assertEqual(b.num, model_b.num, msg=f"bin index {bin_index}")

    self.assertLen(usage_stats.worst_performers, 10)
    for worst_performer, flow_d in zip(
        usage_stats.worst_performers, reversed(flow_data)
    ):
      (
          client_id,
          flow_id,
          (user_cpu_time, system_cpu_time, network_bytes_sent),
      ) = flow_d
      self.assertEqual(
          rdf_client.ClientURN.FromHumanReadable(
              worst_performer.client_id
          ).Basename(),
          client_id,
      )
      self.assertAlmostEqual(
          worst_performer.cpu_usage.user_cpu_time, user_cpu_time
      )
      self.assertAlmostEqual(
          worst_performer.cpu_usage.system_cpu_time, system_cpu_time
      )
      self.assertEqual(worst_performer.network_bytes_sent, network_bytes_sent)
      self.assertEqual(
          rdfvalue.SessionID.FromHumanReadable(
              worst_performer.session_id
          ).Path(),
          "/%s/%s" % (client_id, flow_id),
      )

  def testReadHuntClientResourcesStatsCorrectlyAggregatesVeryLargeNumbers(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id_1 = db_test_utils.InitializeClient(self.db)
    db_test_utils.InitializeFlow(
        self.db,
        client_id_1,
        flow_id=hunt_id,
        flow_state=rdf_flow_objects.Flow.FlowState.FINISHED,
        parent_hunt_id=hunt_id,
        cpu_time_used=rdf_client_stats.CpuSeconds(
            user_cpu_time=3810072130, system_cpu_time=3810072130
        ),
        network_bytes_sent=3810072130,
    )

    client_id_2 = db_test_utils.InitializeClient(self.db)
    db_test_utils.InitializeFlow(
        self.db,
        client_id_2,
        flow_id=hunt_id,
        flow_state=rdf_flow_objects.Flow.FlowState.FINISHED,
        parent_hunt_id=hunt_id,
        cpu_time_used=rdf_client_stats.CpuSeconds(
            user_cpu_time=2143939532, system_cpu_time=2143939532
        ),
        network_bytes_sent=2143939532,
    )

    usage_stats = self.db.ReadHuntClientResourcesStats(hunt_id)

    self.assertEqual(usage_stats.user_cpu_stats.num, 2)
    self.assertAlmostEqual(usage_stats.user_cpu_stats.sum, 5954011662, 5)
    self.assertAlmostEqual(usage_stats.user_cpu_stats.stddev, 833066299, 5)
    self.assertAlmostEqual(usage_stats.system_cpu_stats.sum, 5954011662, 5)
    self.assertAlmostEqual(usage_stats.system_cpu_stats.stddev, 833066299, 5)
    self.assertAlmostEqual(
        usage_stats.network_bytes_sent_stats.sum, 5954011662, 5
    )
    self.assertAlmostEqual(
        usage_stats.network_bytes_sent_stats.stddev, 833066299, 5
    )
    self.assertLen(usage_stats.worst_performers, 2)
    self.assertEqual(
        rdf_client.ClientURN.FromHumanReadable(
            usage_stats.worst_performers[0].client_id
        ).Path(),
        f"/{client_id_1}",
    )
    self.assertAlmostEqual(
        usage_stats.worst_performers[0].cpu_usage.user_cpu_time, 3810072130.0
    )
    self.assertAlmostEqual(
        usage_stats.worst_performers[0].cpu_usage.system_cpu_time, 3810072130.0
    )
    self.assertEqual(
        usage_stats.worst_performers[0].network_bytes_sent, 3810072130
    )
    self.assertEqual(
        rdfvalue.SessionID.FromHumanReadable(
            usage_stats.worst_performers[0].session_id
        ),
        f"/{client_id_1}/{hunt_id}",
    )
    self.assertEqual(
        rdf_client.ClientURN.FromHumanReadable(
            usage_stats.worst_performers[1].client_id
        ).Path(),
        f"/{client_id_2}",
    )
    self.assertAlmostEqual(
        usage_stats.worst_performers[1].cpu_usage.user_cpu_time, 2143939532.0
    )
    self.assertAlmostEqual(
        usage_stats.worst_performers[1].cpu_usage.system_cpu_time, 2143939532.0
    )
    self.assertEqual(
        usage_stats.worst_performers[1].network_bytes_sent, 2143939532
    )
    self.assertEqual(
        rdfvalue.SessionID.FromHumanReadable(
            usage_stats.worst_performers[1].session_id
        ),
        f"/{client_id_2}/{hunt_id}",
    )

  def testReadHuntClientResourcesStatsFiltersDirectFlowIdToMatchTheHuntID(self):
    client_id = db_test_utils.InitializeClient(self.db)
    hunt_id = db_test_utils.InitializeHunt(self.db)

    # The `flow_id` is randomly initialized, so it will not match the `hunt_id`.
    db_test_utils.InitializeFlow(
        self.db,
        client_id,
        parent_hunt_id=hunt_id,
    )

    usage_stats = self.db.ReadHuntClientResourcesStats(hunt_id)
    self.assertEqual(usage_stats.user_cpu_stats.num, 0)
    self.assertEqual(usage_stats.system_cpu_stats.num, 0)
    self.assertEqual(usage_stats.network_bytes_sent_stats.num, 0)
    self.assertEmpty(usage_stats.worst_performers)

  def testReadHuntFlowsStatesAndTimestampsWorksCorrectlyForMultipleFlows(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    expected = []
    for i in range(10):
      client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)

      if i % 2 == 0:
        flow_state = flows_pb2.Flow.FlowState.RUNNING
      else:
        flow_state = flows_pb2.Flow.FlowState.FINISHED
      self.db.UpdateFlow(client_id, flow_id, flow_state=flow_state)

      flow_obj = self.db.ReadFlowObject(client_id, flow_id)
      expected.append(
          db.FlowStateAndTimestamps(
              flow_state=flow_obj.flow_state,
              create_time=flow_obj.create_time,
              last_update_time=flow_obj.last_update_time,
          )
      )

    state_and_times = self.db.ReadHuntFlowsStatesAndTimestamps(hunt_id)
    self.assertCountEqual(state_and_times, expected)

  def testReadHuntFlowsStatesAndTimestampsIgnoresNestedFlows(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)
    self._SetupHuntClientAndFlow(
        hunt_id=hunt_id,
        client_id=client_id,
        flow_id=flow.RandomFlowId(),
        parent_flow_id=flow_id,
    )

    state_and_times = self.db.ReadHuntFlowsStatesAndTimestamps(hunt_id)
    self.assertLen(state_and_times, 1)

    flow_obj = self.db.ReadFlowObject(client_id, flow_id)
    self.assertEqual(
        state_and_times[0],
        db.FlowStateAndTimestamps(
            flow_state=flow_obj.flow_state,
            create_time=flow_obj.create_time,
            last_update_time=flow_obj.last_update_time,
        ),
    )

  def testReadHuntOutputPluginLogEntriesReturnsEntryFromSingleHuntFlow(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    output_plugin_id = "1"
    client_id, flow_id = self._SetupHuntClientAndFlow(
        client_id="C.12345678901234aa", hunt_id=hunt_id
    )
    self.db.WriteFlowOutputPluginLogEntry(
        flows_pb2.FlowOutputPluginLogEntry(
            client_id=client_id,
            flow_id=flow_id,
            output_plugin_id=output_plugin_id,
            hunt_id=hunt_id,
            message="blah",
        )
    )

    hunt_op_log_entries = self.db.ReadHuntOutputPluginLogEntries(
        hunt_id, output_plugin_id, 0, 10
    )
    self.assertLen(hunt_op_log_entries, 1)
    self.assertIsInstance(
        hunt_op_log_entries[0], flows_pb2.FlowOutputPluginLogEntry
    )
    self.assertEqual(hunt_op_log_entries[0].hunt_id, hunt_id)
    self.assertEqual(hunt_op_log_entries[0].client_id, client_id)
    self.assertEqual(hunt_op_log_entries[0].flow_id, flow_id)
    self.assertEqual(hunt_op_log_entries[0].message, "blah")

  def _WriteHuntOutputPluginLogEntries(self) -> str:
    hunt_id = db_test_utils.InitializeHunt(self.db)

    output_plugin_id = "1"
    for i in range(10):
      client_id, flow_id = self._SetupHuntClientAndFlow(
          client_id="C.12345678901234a%d" % i, hunt_id=hunt_id
      )
      enum = flows_pb2.FlowOutputPluginLogEntry.LogEntryType
      if i % 3 == 0:
        log_entry_type = enum.ERROR
      else:
        log_entry_type = enum.LOG
      self.db.WriteFlowOutputPluginLogEntry(
          flows_pb2.FlowOutputPluginLogEntry(
              client_id=client_id,
              flow_id=flow_id,
              hunt_id=hunt_id,
              output_plugin_id=output_plugin_id,
              log_entry_type=log_entry_type,
              message="blah%d" % i,
          )
      )

    return hunt_id

  def testReadHuntOutputPluginLogEntriesReturnsEntryFromMultipleHuntFlows(self):
    hunt_id = self._WriteHuntOutputPluginLogEntries()

    hunt_op_log_entries = self.db.ReadHuntOutputPluginLogEntries(
        hunt_id, "1", 0, 100
    )
    self.assertLen(hunt_op_log_entries, 10)
    # Make sure messages are returned in timestamps-ascending order.
    for i, e in enumerate(hunt_op_log_entries):
      self.assertEqual(e.message, "blah%d" % i)

  def testReadHuntOutputPluginLogEntriesCorrectlyAppliesOffsetAndCountFilters(
      self,
  ):
    hunt_id = self._WriteHuntOutputPluginLogEntries()

    for i in range(10):
      hunt_op_log_entries = self.db.ReadHuntOutputPluginLogEntries(
          hunt_id, "1", i, 1
      )
      self.assertLen(hunt_op_log_entries, 1)
      self.assertEqual(hunt_op_log_entries[0].message, "blah%d" % i)

  def testReadHuntOutputPluginLogEntriesCorrectlyAppliesWithTypeFilter(self):
    hunt_id = self._WriteHuntOutputPluginLogEntries()

    hunt_op_log_entries = self.db.ReadHuntOutputPluginLogEntries(
        hunt_id,
        "1",
        0,
        100,
        with_type=flows_pb2.FlowOutputPluginLogEntry.LogEntryType.UNSET,
    )
    self.assertEmpty(hunt_op_log_entries)

    hunt_op_log_entries = self.db.ReadHuntOutputPluginLogEntries(
        hunt_id,
        "1",
        0,
        100,
        with_type=flows_pb2.FlowOutputPluginLogEntry.LogEntryType.ERROR,
    )
    self.assertLen(hunt_op_log_entries, 4)

    hunt_op_log_entries = self.db.ReadHuntOutputPluginLogEntries(
        hunt_id,
        "1",
        0,
        100,
        with_type=flows_pb2.FlowOutputPluginLogEntry.LogEntryType.LOG,
    )
    self.assertLen(hunt_op_log_entries, 6)

  def testReadHuntOutputPluginLogEntriesCorrectlyAppliesCombinationOfFilters(
      self,
  ):
    hunt_id = self._WriteHuntOutputPluginLogEntries()

    hunt_log_entries = self.db.ReadHuntOutputPluginLogEntries(
        hunt_id,
        "1",
        0,
        1,
        with_type=flows_pb2.FlowOutputPluginLogEntry.LogEntryType.LOG,
    )
    self.assertLen(hunt_log_entries, 1)
    self.assertEqual(hunt_log_entries[0].message, "blah1")

  def testCountHuntOutputPluginLogEntriesReturnsCorrectCount(self):
    hunt_id = self._WriteHuntOutputPluginLogEntries()

    num_entries = self.db.CountHuntOutputPluginLogEntries(hunt_id, "1")
    self.assertEqual(num_entries, 10)

  def testCountHuntOutputPluginLogEntriesRespectsWithTypeFilter(self):
    hunt_id = self._WriteHuntOutputPluginLogEntries()

    num_entries = self.db.CountHuntOutputPluginLogEntries(
        hunt_id,
        "1",
        with_type=flows_pb2.FlowOutputPluginLogEntry.LogEntryType.LOG,
    )
    self.assertEqual(num_entries, 6)

    num_entries = self.db.CountHuntOutputPluginLogEntries(
        hunt_id,
        "1",
        with_type=flows_pb2.FlowOutputPluginLogEntry.LogEntryType.ERROR,
    )
    self.assertEqual(num_entries, 4)

  def testFlowStateUpdateUsingUpdateFlow(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(
        hunt_id=hunt_id, flow_state=rdf_flow_objects.Flow.FlowState.RUNNING
    )

    results = self.db.ReadHuntFlows(
        hunt_id,
        0,
        10,
        filter_condition=db.HuntFlowsCondition.FLOWS_IN_PROGRESS_ONLY,
    )
    self.assertLen(results, 1)

    results = self.db.ReadHuntFlows(
        hunt_id,
        0,
        10,
        filter_condition=db.HuntFlowsCondition.COMPLETED_FLOWS_ONLY,
    )
    self.assertEmpty(results)

    proto_flow = self.db.ReadFlowObject(client_id, flow_id)
    proto_flow.flow_state = flows_pb2.Flow.FlowState.FINISHED
    self.db.UpdateFlow(client_id, proto_flow.flow_id, flow_obj=proto_flow)

    results = self.db.ReadHuntFlows(
        hunt_id,
        0,
        10,
        filter_condition=db.HuntFlowsCondition.FLOWS_IN_PROGRESS_ONLY,
    )
    self.assertEmpty(results)

    results = self.db.ReadHuntFlows(
        hunt_id,
        0,
        10,
        filter_condition=db.HuntFlowsCondition.COMPLETED_FLOWS_ONLY,
    )
    self.assertLen(results, 1)

  def testFlowStateUpdateUsingReleaseProcessedFlow(self):
    hunt_id = db_test_utils.InitializeHunt(self.db)

    client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_id)

    flow_obj = self.db.LeaseFlowForProcessing(
        client_id, flow_id, rdfvalue.Duration.From(1, rdfvalue.MINUTES)
    )
    self.assertEqual(flow_obj.flow_state, flows_pb2.Flow.FlowState.UNSET)

    flow_obj.flow_state = flows_pb2.Flow.FlowState.ERROR
    self.db.ReleaseProcessedFlow(flow_obj)

    results = self.db.ReadHuntFlows(
        hunt_id, 0, 10, filter_condition=db.HuntFlowsCondition.FAILED_FLOWS_ONLY
    )
    self.assertLen(results, 1)


# This file is a test library and thus does not require a __main__ block.
