#!/usr/bin/env python
"""Test for the flow state class."""

from absl import app

from grr_response_core.lib import rdfvalue
from grr_response_core.lib.rdfvalues import test_base as rdf_test_base
from grr.test_lib import test_lib


class SessionIDTest(rdf_test_base.RDFValueTestMixin, test_lib.GRRBaseTest):
  """Test SessionID."""

  rdfvalue_class = rdfvalue.SessionID

  def GenerateSample(self, number=0):
    id_str = "%08X" % (number % 2**32)
    return rdfvalue.SessionID(flow_name=id_str)

  def testSessionIDValidation(self):
    rdfvalue.SessionID(rdfvalue.RDFURN("aff4:/flows/A:12345678"))
    rdfvalue.SessionID(rdfvalue.RDFURN("aff4:/flows/A:TransferStore"))
    rdfvalue.SessionID(rdfvalue.RDFURN("aff4:/flows/DEBUG-user1:12345678"))
    rdfvalue.SessionID(rdfvalue.RDFURN("aff4:/flows/DEBUG-user1:12345678:hunt"))

  def testBadStructure(self):
    self.assertRaises(
        rdfvalue.InitializeError,
        rdfvalue.SessionID,
        rdfvalue.RDFURN("aff4:/flows/A:123456:1:"),
    )
    self.assertRaises(
        rdfvalue.InitializeError,
        rdfvalue.SessionID,
        rdfvalue.RDFURN("aff4:/flows/A:123456::"),
    )
    self.assertRaises(
        rdfvalue.InitializeError,
        rdfvalue.SessionID,
        rdfvalue.RDFURN("aff4:/flows/A:123456:"),
    )
    self.assertRaises(
        rdfvalue.InitializeError,
        rdfvalue.SessionID,
        rdfvalue.RDFURN("aff4:/flows/A:"),
    )
    self.assertRaises(
        rdfvalue.InitializeError,
        rdfvalue.SessionID,
        rdfvalue.RDFURN("aff4:/flows/:"),
    )

  def testBadQueue(self):
    self.assertRaises(
        rdfvalue.InitializeError,
        rdfvalue.SessionID,
        rdfvalue.RDFURN("aff4:/flows/A%b:12345678"),
    )

  def testBadFlowID(self):
    self.assertRaises(
        rdfvalue.InitializeError,
        rdfvalue.SessionID,
        rdfvalue.RDFURN("aff4:/flows/A:1234567G%sdf"),
    )


def main(argv):
  # Run the full test suite
  test_lib.main(argv)


if __name__ == "__main__":
  app.run(main)