ray/test/test.py

from __future__ import print_function

import os
import signal
import socket
import subprocess
import sys
import unittest
import random
import time
import tempfile

import plasma

USE_VALGRIND = False

def random_object_id():
  return "".join([chr(random.randint(0, 255)) for _ in range(20)])

def generate_metadata(length):
  metadata = length * ["\x00"]
  if length > 0:
    metadata[0] = chr(random.randint(0, 255))
    metadata[-1] = chr(random.randint(0, 255))
    for _ in range(100):
      metadata[random.randint(0, length - 1)] = chr(random.randint(0, 255))
  return buffer("".join(metadata))

def write_to_data_buffer(buff, length):
  if length > 0:
    buff[0] = chr(random.randint(0, 255))
    buff[-1] = chr(random.randint(0, 255))
    for _ in range(100):
      buff[random.randint(0, length - 1)] = chr(random.randint(0, 255))

def create_object(client, data_size, metadata_size, seal=True):
  object_id = random_object_id()
  metadata = generate_metadata(metadata_size)
  memory_buffer = client.create(object_id, data_size, metadata)
  write_to_data_buffer(memory_buffer, data_size)
  if seal:
    client.seal(object_id)
  return object_id, memory_buffer, metadata

class TestPlasmaClient(unittest.TestCase):

  def setUp(self):
    # Start Plasma.
    plasma_store_executable = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../build/plasma_store")
    store_name = "/tmp/store{}".format(random.randint(0, 10000))
    command = [plasma_store_executable, "-s", store_name]
    if USE_VALGRIND:
      self.p = subprocess.Popen(["valgrind", "--track-origins=yes", "--leak-check=full"] + command)
      time.sleep(2.0)
    else:
      self.p = subprocess.Popen(command)
    # Connect to Plasma.
    self.plasma_client = plasma.PlasmaClient(store_name)

  def tearDown(self):
    # Kill the plasma store process.
    if USE_VALGRIND:
      self.p.send_signal(signal.SIGTERM)
      self.p.wait()
      if self.p.returncode != 0:
        os._exit(-1)
    else:
      self.p.kill()

  def test_create(self):
    # Create an object id string.
    object_id = random_object_id()
    # Create a new buffer and write to it.
    length = 50
    memory_buffer = self.plasma_client.create(object_id, length)
    for i in range(length):
      memory_buffer[i] = chr(i % 256)
    # Seal the object.
    self.plasma_client.seal(object_id)
    # Get the object.
    memory_buffer = self.plasma_client.get(object_id)
    for i in range(length):
      self.assertEqual(memory_buffer[i], chr(i % 256))

  def test_create_with_metadata(self):
    for length in range(1000):
      # Create an object id string.
      object_id = random_object_id()
      # Create a random metadata string.
      metadata = generate_metadata(length)
      # Create a new buffer and write to it.
      memory_buffer = self.plasma_client.create(object_id, length, metadata)
      for i in range(length):
        memory_buffer[i] = chr(i % 256)
      # Seal the object.
      self.plasma_client.seal(object_id)
      # Get the object.
      memory_buffer = self.plasma_client.get(object_id)
      for i in range(length):
        self.assertEqual(memory_buffer[i], chr(i % 256))
      # Get the metadata.
      metadata_buffer = self.plasma_client.get_metadata(object_id)
      self.assertEqual(len(metadata), len(metadata_buffer))
      for i in range(len(metadata)):
        self.assertEqual(metadata[i], metadata_buffer[i])

  def test_contains(self):
    fake_object_ids = [random_object_id() for _ in range(100)]
    real_object_ids = [random_object_id() for _ in range(100)]
    for object_id in real_object_ids:
      self.assertFalse(self.plasma_client.contains(object_id))
      memory_buffer = self.plasma_client.create(object_id, 100)
      self.plasma_client.seal(object_id)
      self.assertTrue(self.plasma_client.contains(object_id))
    for object_id in fake_object_ids:
      self.assertFalse(self.plasma_client.contains(object_id))
    for object_id in real_object_ids:
      self.assertTrue(self.plasma_client.contains(object_id))

  def test_individual_delete(self):
    length = 100
    # Create an object id string.
    object_id = random_object_id()
    # Create a random metadata string.
    metadata = generate_metadata(100)
    # Create a new buffer and write to it.
    memory_buffer = self.plasma_client.create(object_id, length, metadata)
    for i in range(length):
      memory_buffer[i] = chr(i % 256)
    # Seal the object.
    self.plasma_client.seal(object_id)
    # Check that the object is present.
    self.assertTrue(self.plasma_client.contains(object_id))
    # Delete the object.
    self.plasma_client.delete(object_id)
    # Make sure the object is no longer present.
    self.assertFalse(self.plasma_client.contains(object_id))

  def test_delete(self):
    # Create some objects.
    object_ids = [random_object_id() for _ in range(100)]
    for object_id in object_ids:
      length = 100
      # Create a random metadata string.
      metadata = generate_metadata(100)
      # Create a new buffer and write to it.
      memory_buffer = self.plasma_client.create(object_id, length, metadata)
      for i in range(length):
        memory_buffer[i] = chr(i % 256)
      # Seal the object.
      self.plasma_client.seal(object_id)
      # Check that the object is present.
      self.assertTrue(self.plasma_client.contains(object_id))

    # Delete the objects and make sure they are no longer present.
    for object_id in object_ids:
      # Delete the object.
      self.plasma_client.delete(object_id)
      # Make sure the object is no longer present.
      self.assertFalse(self.plasma_client.contains(object_id))

  def test_illegal_functionality(self):
    # Create an object id string.
    object_id = random_object_id()
    # Create a new buffer and write to it.
    length = 1000
    memory_buffer = self.plasma_client.create(object_id, length)
    # Make sure we cannot access memory out of bounds.
    self.assertRaises(Exception, lambda : memory_buffer[length])
    # Seal the object.
    self.plasma_client.seal(object_id)
    # This test is commented out because it currently fails.
    # # Make sure the object is ready only now.
    # def illegal_assignment():
    #   memory_buffer[0] = chr(0)
    # self.assertRaises(Exception, illegal_assignment)
    # Get the object.
    memory_buffer = self.plasma_client.get(object_id)
    # Make sure the object is read only.
    def illegal_assignment():
      memory_buffer[0] = chr(0)
    self.assertRaises(Exception, illegal_assignment)

class TestPlasmaManager(unittest.TestCase):

  def setUp(self):
    # Start two PlasmaStores.
    plasma_store_executable = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../build/plasma_store")
    store_name1 = "/tmp/store{}".format(random.randint(0, 10000))
    store_name2 = "/tmp/store{}".format(random.randint(0, 10000))
    plasma_store_command1 = [plasma_store_executable, "-s", store_name1]
    plasma_store_command2 = [plasma_store_executable, "-s", store_name2]

    if USE_VALGRIND:
      self.p2 = subprocess.Popen(["valgrind", "--track-origins=yes", "--error-exitcode=1"] + plasma_store_command1)
      self.p3 = subprocess.Popen(["valgrind", "--track-origins=yes", "--error-exitcode=1"] + plasma_store_command2)
    else:
      self.p2 = subprocess.Popen(plasma_store_command1)
      self.p3 = subprocess.Popen(plasma_store_command2)

    # Start two PlasmaManagers.
    self.port1 = random.randint(10000, 50000)
    self.port2 = random.randint(10000, 50000)
    plasma_manager_executable = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../build/plasma_manager")
    plasma_manager_command1 = [plasma_manager_executable, "-s", store_name1, "-m", "127.0.0.1", "-p", str(self.port1)]
    plasma_manager_command2 = [plasma_manager_executable, "-s", store_name2, "-m", "127.0.0.1", "-p", str(self.port2)]

    if USE_VALGRIND:
      self.p4 = subprocess.Popen(["valgrind", "--track-origins=yes", "--error-exitcode=1"] + plasma_manager_command1)
      self.p5 = subprocess.Popen(["valgrind", "--track-origins=yes", "--error-exitcode=1"] + plasma_manager_command2)
      time.sleep(2.0)
    else:
      self.p4 = subprocess.Popen(plasma_manager_command1)
      self.p5 = subprocess.Popen(plasma_manager_command2)
      time.sleep(0.1)

    # Connect two PlasmaClients.
    self.client1 = plasma.PlasmaClient(store_name1, "127.0.0.1", self.port1)
    self.client2 = plasma.PlasmaClient(store_name2, "127.0.0.1", self.port2)
    time.sleep(0.5)

  def tearDown(self):
    # Kill the PlasmaStore and PlasmaManager processes.
    if USE_VALGRIND:
      self.p4.send_signal(signal.SIGTERM)
      self.p4.wait()
      self.p5.send_signal(signal.SIGTERM)
      self.p5.wait()
      self.p2.send_signal(signal.SIGTERM)
      self.p2.wait()
      self.p3.send_signal(signal.SIGTERM)
      self.p3.wait()
      if self.p2.returncode != 0 or self.p3.returncode != 0 or self.p4.returncode != 0 or self.p5.returncode != 0:
        print("aborting due to valgrind error")
        os._exit(-1)
    else:
      self.p2.kill()
      self.p3.kill()
      self.p4.kill()
      self.p5.kill()

  def test_transfer(self):
    for _ in range(100):
      # Create an object.
      object_id1, memory_buffer1, metadata1 = create_object(self.client1, 2000, 2000)
      # Transfer the buffer to the the other PlasmaStore.
      self.client1.transfer("127.0.0.1", self.port2, object_id1)
      # Compare the two buffers.
      self.assertEqual(memory_buffer1[:], self.client2.get(object_id1)[:])
      self.assertEqual(self.client1.get(object_id1)[:], self.client2.get(object_id1)[:])
      self.assertEqual(metadata1[:], self.client2.get_metadata(object_id1)[:])
      self.assertEqual(self.client1.get_metadata(object_id1)[:], self.client2.get_metadata(object_id1)[:])
      # Transfer the buffer again.
      self.client1.transfer("127.0.0.1", self.port2, object_id1)
      self.assertEqual(metadata1[:], self.client2.get_metadata(object_id1)[:])
      # Compare the two buffers.
      self.assertEqual(self.client1.get(object_id1)[:], self.client2.get(object_id1)[:])

      # Create an object.
      object_id2, memory_buffer2, metadata2 = create_object(self.client2, 20000, 20000)
      # Transfer the buffer to the the other PlasmaStore.
      self.client2.transfer("127.0.0.1", self.port1, object_id2)
      # Compare the two buffers.
      self.assertEqual(memory_buffer2[:], self.client2.get(object_id2)[:])
      self.assertEqual(self.client1.get(object_id2)[:], self.client2.get(object_id2)[:])
      self.assertEqual(metadata2[:], self.client2.get_metadata(object_id2)[:])
      self.assertEqual(self.client1.get_metadata(object_id2)[:], self.client2.get_metadata(object_id2)[:])

  def test_illegal_functionality(self):
    # Create an object id string.
    object_id = random_object_id()
    # Create a new buffer.
    # memory_buffer = self.client1.create(object_id, 20000)
    # This test is commented out because it currently fails.
    # # Transferring the buffer before sealing it should fail.
    # self.assertRaises(Exception, lambda : self.manager1.transfer(1, object_id))

  def test_stresstest(self):
    a = time.time()
    object_ids = []
    for i in range(10000): # TODO(pcm): increase this to 100000
      object_id = random_object_id()
      object_ids.append(object_id)
      self.client1.create(object_id, 1)
      self.client1.seal(object_id)
    for object_id in object_ids:
      self.client1.transfer("127.0.0.1", self.port2, object_id)
    b = time.time() - a

    print("it took", b, "seconds to put and transfer the objects")

if __name__ == "__main__":
  if len(sys.argv) > 1:
    # pop the argument so we don't mess with unittest's own argument parser
    arg = sys.argv.pop()
    if arg == "valgrind":
      USE_VALGRIND = True
      print("Using valgrind for tests")
  unittest.main(verbosity=2)