services/nodemanager/arvnodeman/test/fake_driver.py

   1 import re
   2 import urllib
   3 import ssl
   4 import time
   5
   6 from arvnodeman.computenode import ARVADOS_TIMEFMT
   7
   8 from libcloud.compute.base import NodeSize, Node, NodeDriver, NodeState
   9 from libcloud.common.exceptions import BaseHTTPError
  10
  11 all_nodes = []
  12 create_calls = 0
  13 quota = 2
  14
  15 class FakeDriver(NodeDriver):
  16     def __init__(self, *args, **kwargs):
  17         self.name = "FakeDriver"
  18
  19     def list_sizes(self, **kwargs):
  20         return [NodeSize("Standard_D3", "Standard_D3", 3500, 200, 0, 0, self),
  21                 NodeSize("Standard_D4", "Standard_D4", 7000, 400, 0, 0, self)]
  22
  23     def list_nodes(self, **kwargs):
  24         return all_nodes
  25
  26     def create_node(self, name=None,
  27                     size=None,
  28                     image=None,
  29                     auth=None,
  30                     ex_storage_account=None,
  31                     ex_customdata=None,
  32                     ex_resource_group=None,
  33                     ex_user_name=None,
  34                     ex_tags=None,
  35                     ex_network=None,
  36                     ex_userdata=None):
  37         global all_nodes, create_calls
  38         create_calls += 1
  39         n = Node(name, name, NodeState.RUNNING, [], [], self, size=size, extra={"tags": ex_tags})
  40         all_nodes.append(n)
  41         if ex_customdata:
  42             ping_url = re.search(r"echo '(.*)' > /var/tmp/arv-node-data/arv-ping-url", ex_customdata).groups(1)[0] + "&instance_id=" + name
  43         if ex_userdata:
  44             ping_url = ex_userdata
  45         ctx = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
  46         ctx.verify_mode = ssl.CERT_NONE
  47         f = urllib.urlopen(ping_url, "", context=ctx)
  48         f.close()
  49         return n
  50
  51     def destroy_node(self, cloud_node):
  52         global all_nodes
  53         all_nodes = [n for n in all_nodes if n.id != cloud_node.id]
  54         return True
  55
  56     def get_image(self, img):
  57         pass
  58
  59     def ex_create_tags(self, cloud_node, tags):
  60         pass
  61
  62 class QuotaDriver(FakeDriver):
  63     def create_node(self, name=None,
  64                     size=None,
  65                     image=None,
  66                     auth=None,
  67                     ex_storage_account=None,
  68                     ex_customdata=None,
  69                     ex_resource_group=None,
  70                     ex_user_name=None,
  71                     ex_tags=None,
  72                     ex_network=None):
  73         global all_nodes, create_calls, quota
  74         if len(all_nodes) >= quota:
  75             raise BaseHTTPError(503, "Quota exceeded")
  76         else:
  77             return super(QuotaDriver, self).create_node(name=name,
  78                     size=size,
  79                     image=image,
  80                     auth=auth,
  81                     ex_storage_account=ex_storage_account,
  82                     ex_customdata=ex_customdata,
  83                     ex_resource_group=ex_resource_group,
  84                     ex_user_name=ex_user_name,
  85                     ex_tags=ex_tags,
  86                     ex_network=ex_network)
  87
  88     def destroy_node(self, cloud_node):
  89         global all_nodes, quota
  90         all_nodes = [n for n in all_nodes if n.id != cloud_node.id]
  91         if len(all_nodes) == 0:
  92             quota = 4
  93         return True
  94
  95 class FailingDriver(FakeDriver):
  96     def create_node(self, name=None,
  97                     size=None,
  98                     image=None,
  99                     auth=None,
 100                     ex_storage_account=None,
 101                     ex_customdata=None,
 102                     ex_resource_group=None,
 103                     ex_user_name=None,
 104                     ex_tags=None,
 105                     ex_network=None):
 106         raise Exception("nope")
 107
 108 class RetryDriver(FakeDriver):
 109     def create_node(self, name=None,
 110                     size=None,
 111                     image=None,
 112                     auth=None,
 113                     ex_storage_account=None,
 114                     ex_customdata=None,
 115                     ex_resource_group=None,
 116                     ex_user_name=None,
 117                     ex_tags=None,
 118                     ex_network=None):
 119         global create_calls
 120         create_calls += 1
 121         if create_calls < 2:
 122             raise BaseHTTPError(429, "Rate limit exceeded",
 123                                 {'retry-after': '12'})
 124         else:
 125             return super(RetryDriver, self).create_node(name=name,
 126                     size=size,
 127                     image=image,
 128                     auth=auth,
 129                     ex_storage_account=ex_storage_account,
 130                     ex_customdata=ex_customdata,
 131                     ex_resource_group=ex_resource_group,
 132                     ex_user_name=ex_user_name,
 133                     ex_tags=ex_tags,
 134                     ex_network=ex_network)
 135
 136 class FakeAwsDriver(FakeDriver):
 137
 138     def create_node(self, name=None,
 139                     size=None,
 140                     image=None,
 141                     auth=None,
 142                     ex_userdata=None,
 143                     ex_blockdevicemappings=None):
 144         n = super(FakeAwsDriver, self).create_node(name=name,
 145                                                       size=size,
 146                                                       image=image,
 147                                                       auth=auth,
 148                                                       ex_userdata=ex_userdata)
 149         n.extra = {"launch_time": time.strftime(ARVADOS_TIMEFMT, time.gmtime())[:-1]}
 150         return n
 151
 152     def list_sizes(self, **kwargs):
 153         return [NodeSize("m3.xlarge", "Extra Large Instance", 3500, 80, 0, 0, self),
 154                 NodeSize("m4.xlarge", "Extra Large Instance", 3500, 0, 0, 0, self),
 155                 NodeSize("m4.2xlarge", "Double Extra Large Instance", 7000, 0, 0, 0, self)]