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, NodeImage
   9 from libcloud.compute.drivers.gce import GCEDiskType
  10 from libcloud.common.exceptions import BaseHTTPError
  11
  12 all_nodes = []
  13 create_calls = 0
  14 quota = 2
  15
  16 class FakeDriver(NodeDriver):
  17     def __init__(self, *args, **kwargs):
  18         self.name = "FakeDriver"
  19
  20     def list_sizes(self, **kwargs):
  21         return [NodeSize("Standard_D3", "Standard_D3", 3500, 200, 0, 0, self),
  22                 NodeSize("Standard_D4", "Standard_D4", 7000, 400, 0, 0, self)]
  23
  24     def list_nodes(self, **kwargs):
  25         return all_nodes
  26
  27     def create_node(self, name=None,
  28                     size=None,
  29                     image=None,
  30                     auth=None,
  31                     ex_storage_account=None,
  32                     ex_customdata=None,
  33                     ex_resource_group=None,
  34                     ex_user_name=None,
  35                     ex_tags=None,
  36                     ex_metadata=None,
  37                     ex_network=None,
  38                     ex_userdata=None):
  39         global all_nodes, create_calls
  40         create_calls += 1
  41         nodeid = "node%i" % create_calls
  42         n = Node(nodeid, nodeid, NodeState.RUNNING, [], [], self, size=size, extra={"tags": ex_tags})
  43         all_nodes.append(n)
  44         if ex_customdata:
  45             ping_url = re.search(r"echo '(.*)' > /var/tmp/arv-node-data/arv-ping-url", ex_customdata).groups(1)[0]
  46         if ex_userdata:
  47             ping_url = ex_userdata
  48         if ex_metadata:
  49             ping_url = ex_metadata["arv-ping-url"]
  50         ping_url += "&instance_id=" + nodeid
  51         ctx = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
  52         ctx.verify_mode = ssl.CERT_NONE
  53         f = urllib.urlopen(ping_url, "", context=ctx)
  54         f.close()
  55         return n
  56
  57     def destroy_node(self, cloud_node):
  58         global all_nodes
  59         all_nodes = [n for n in all_nodes if n.id != cloud_node.id]
  60         return True
  61
  62     def get_image(self, img):
  63         pass
  64
  65     def ex_create_tags(self, cloud_node, tags):
  66         pass
  67
  68 class QuotaDriver(FakeDriver):
  69     def create_node(self, name=None,
  70                     size=None,
  71                     image=None,
  72                     auth=None,
  73                     ex_storage_account=None,
  74                     ex_customdata=None,
  75                     ex_resource_group=None,
  76                     ex_user_name=None,
  77                     ex_tags=None,
  78                     ex_network=None):
  79         global all_nodes, create_calls, quota
  80         if len(all_nodes) >= quota:
  81             raise BaseHTTPError(503, "Quota exceeded")
  82         else:
  83             return super(QuotaDriver, self).create_node(name=name,
  84                     size=size,
  85                     image=image,
  86                     auth=auth,
  87                     ex_storage_account=ex_storage_account,
  88                     ex_customdata=ex_customdata,
  89                     ex_resource_group=ex_resource_group,
  90                     ex_user_name=ex_user_name,
  91                     ex_tags=ex_tags,
  92                     ex_network=ex_network)
  93
  94     def destroy_node(self, cloud_node):
  95         global all_nodes, quota
  96         all_nodes = [n for n in all_nodes if n.id != cloud_node.id]
  97         if len(all_nodes) == 0:
  98             quota = 4
  99         return True
 100
 101 class FailingDriver(FakeDriver):
 102     def create_node(self, name=None,
 103                     size=None,
 104                     image=None,
 105                     auth=None,
 106                     ex_storage_account=None,
 107                     ex_customdata=None,
 108                     ex_resource_group=None,
 109                     ex_user_name=None,
 110                     ex_tags=None,
 111                     ex_network=None):
 112         raise Exception("nope")
 113
 114 class RetryDriver(FakeDriver):
 115     def create_node(self, name=None,
 116                     size=None,
 117                     image=None,
 118                     auth=None,
 119                     ex_storage_account=None,
 120                     ex_customdata=None,
 121                     ex_resource_group=None,
 122                     ex_user_name=None,
 123                     ex_tags=None,
 124                     ex_network=None):
 125         global create_calls
 126         create_calls += 1
 127         if create_calls < 2:
 128             raise BaseHTTPError(429, "Rate limit exceeded",
 129                                 {'retry-after': '12'})
 130         else:
 131             return super(RetryDriver, self).create_node(name=name,
 132                     size=size,
 133                     image=image,
 134                     auth=auth,
 135                     ex_storage_account=ex_storage_account,
 136                     ex_customdata=ex_customdata,
 137                     ex_resource_group=ex_resource_group,
 138                     ex_user_name=ex_user_name,
 139                     ex_tags=ex_tags,
 140                     ex_network=ex_network)
 141
 142 class FakeAwsDriver(FakeDriver):
 143
 144     def create_node(self, name=None,
 145                     size=None,
 146                     image=None,
 147                     auth=None,
 148                     ex_userdata=None,
 149                     ex_blockdevicemappings=None):
 150         n = super(FakeAwsDriver, self).create_node(name=name,
 151                                                       size=size,
 152                                                       image=image,
 153                                                       auth=auth,
 154                                                       ex_userdata=ex_userdata)
 155         n.extra = {"launch_time": time.strftime(ARVADOS_TIMEFMT, time.gmtime())[:-1]}
 156         return n
 157
 158     def list_sizes(self, **kwargs):
 159         return [NodeSize("m3.xlarge", "Extra Large Instance", 3500, 80, 0, 0, self),
 160                 NodeSize("m4.xlarge", "Extra Large Instance", 3500, 0, 0, 0, self),
 161                 NodeSize("m4.2xlarge", "Double Extra Large Instance", 7000, 0, 0, 0, self)]
 162
 163
 164 class FakeGceDriver(FakeDriver):
 165
 166     def create_node(self, name=None,
 167                     size=None,
 168                     image=None,
 169                     auth=None,
 170                     external_ip=None,
 171                     ex_metadata=None,
 172                     ex_tags=None,
 173                     ex_disks_gce_struct=None):
 174         n = super(FakeGceDriver, self).create_node(name=name,
 175                                                    size=size,
 176                                                    image=image,
 177                                                    auth=auth,
 178                                                    ex_metadata=ex_metadata)
 179         n.extra = {
 180             "metadata": {
 181                 "items": [{"key": k, "value": v} for k,v in ex_metadata.iteritems()]
 182             },
 183             "zone": "fake"
 184         }
 185         return n
 186
 187     def list_images(self, ex_project=None):
 188         return [NodeImage("fake_image_id", "fake_image_id", self)]
 189
 190     def list_sizes(self, **kwargs):
 191         return [NodeSize("n1-standard-1", "Standard", 3750, None, 0, 0, self),
 192                 NodeSize("n1-standard-2", "Double standard", 7500, None, 0, 0, self)]
 193
 194     def ex_list_disktypes(self, zone=None):
 195         return [GCEDiskType("pd-standard", "pd-standard", zone, self,
 196                             extra={"selfLink": "pd-standard"}),
 197                 GCEDiskType("local-ssd", "local-ssd", zone, self,
 198                             extra={"selfLink": "local-ssd"})]
 199
 200     def ex_get_node(self, name, zone=None):
 201         global all_nodes
 202         for n in all_nodes:
 203             if n.id == name:
 204                 return n
 205         return None
 206
 207     def ex_set_node_metadata(self, n, items):
 208         n.extra["metadata"]["items"] = items