github.com/google/fleetspeak@v0.1.15-0.20240426164851-4f31f62c1aea/fleetspeak_python/fleetspeak/client_connector/testing/testclient.py (about)

     1  #!/usr/bin/env python
     2  # Copyright 2017 Google Inc.
     3  #
     4  # Licensed under the Apache License, Version 2.0 (the "License");
     5  # you may not use this file except in compliance with the License.
     6  # You may obtain a copy of the License at
     7  #
     8  #     https://www.apache.org/licenses/LICENSE-2.0
     9  #
    10  # Unless required by applicable law or agreed to in writing, software
    11  # distributed under the License is distributed on an "AS IS" BASIS,
    12  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  # See the License for the specific language governing permissions and
    14  # limitations under the License.
    15  
    16  """Simplified python testclient for Fleetspeak daemonservice.
    17  
    18  This is a minimal daemonservice client. It is started and exercised by
    19  daemonservice_test.go in order to verify that the python client library
    20  correctly implements the protocol expected by daemonservice.
    21  """
    22  
    23  import logging
    24  import time
    25  
    26  from absl import app
    27  from absl import flags
    28  from fleetspeak.client_connector import connector
    29  
    30  FLAGS = flags.FLAGS
    31  flags.DEFINE_string(
    32      "mode", "loopback", "Mode of operation. Options are: loopback"
    33  )
    34  
    35  
    36  class FatalError(Exception):
    37    pass
    38  
    39  
    40  def Loopback():
    41    logging.info("starting loopback")
    42    con = connector.FleetspeakConnection(version="0.5")
    43    logging.info("connection created")
    44    while True:
    45      msg, _ = con.Recv()
    46      msg.message_type += "Response"
    47      con.Send(msg)
    48  
    49  
    50  def MemoryHog():
    51    """Takes 20MB of memory, then sleeps forever."""
    52    logging.info("starting memory leak")
    53    con = connector.FleetspeakConnection(version="0.5")
    54    logging.info("connection created")
    55    buf = "a" * (1024 * 1024 * 20)
    56    while True:
    57      time.sleep(1)
    58  
    59  
    60  def Freezed():
    61    """Connects to Fleetspeak, then sleeps indefinitely."""
    62    logging.info("starting freezed")
    63    con = connector.FleetspeakConnection(version="0.5")
    64    logging.info("connection created")
    65    while True:
    66      time.sleep(1)
    67  
    68  
    69  def Heartbeat():
    70    """Sends a heartbeat every 100ms, indefinitely."""
    71    logging.info("starting heartbeat")
    72    con = connector.FleetspeakConnection(version="0.5")
    73    logging.info("connection created")
    74    while True:
    75      time.sleep(0.1)
    76      con.Heartbeat()
    77  
    78  
    79  def main(argv=None):
    80    del argv  # Unused.
    81  
    82    if FLAGS.mode == "loopback":
    83      Loopback()
    84      return
    85  
    86    if FLAGS.mode == "memoryhog":
    87      MemoryHog()
    88      return
    89  
    90    if FLAGS.mode == "freezed":
    91      Freezed()
    92      return
    93  
    94    if FLAGS.mode == "100ms-heartbeat":
    95      Heartbeat()
    96      return
    97  
    98    raise FatalError("Unknown mode: %s", FLAGS.mode)
    99  
   100  
   101  if __name__ == "__main__":
   102    app.run(main)