X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/74e03b2190fea52b3ea5800c50eca55b6e50b220..76f6b358cf612e67b586452236394feabc464373:/services/nodemanager/arvnodeman/test/fake_driver.py

diff --git a/services/nodemanager/arvnodeman/test/fake_driver.py b/services/nodemanager/arvnodeman/test/fake_driver.py
index 9cb582c953..11ba271b67 100644
--- a/services/nodemanager/arvnodeman/test/fake_driver.py
+++ b/services/nodemanager/arvnodeman/test/fake_driver.py
@@ -11,7 +11,7 @@ from arvnodeman.computenode import ARVADOS_TIMEFMT
 
 from libcloud.compute.base import NodeSize, Node, NodeDriver, NodeState, NodeImage
 from libcloud.compute.drivers.gce import GCEDiskType
-from libcloud.common.exceptions import BaseHTTPError
+from libcloud.common.exceptions import BaseHTTPError, RateLimitReachedError
 
 all_nodes = []
 create_calls = 0
@@ -43,13 +43,16 @@ class FakeDriver(NodeDriver):
         global all_nodes, create_calls
         create_calls += 1
         nodeid = "node%i" % create_calls
+        if ex_tags is None:
+            ex_tags = {}
+        ex_tags.update({'arvados_node_size': size.id})
         n = Node(nodeid, nodeid, NodeState.RUNNING, [], [], self, size=size, extra={"tags": ex_tags})
         all_nodes.append(n)
         if ex_customdata:
             ping_url = re.search(r"echo '(.*)' > /var/tmp/arv-node-data/arv-ping-url", ex_customdata).groups(1)[0]
         if ex_userdata:
             ping_url = ex_userdata
-        if ex_metadata:
+        elif ex_metadata:
             ping_url = ex_metadata["arv-ping-url"]
         ping_url += "&instance_id=" + nodeid
         ctx = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
@@ -129,8 +132,11 @@ class RetryDriver(FakeDriver):
         global create_calls
         create_calls += 1
         if create_calls < 2:
+            raise RateLimitReachedError(429, "Rate limit exceeded",
+                                        headers={'retry-after': '12'})
+        elif create_calls < 3:
             raise BaseHTTPError(429, "Rate limit exceeded",
-                                {'retry-after': '12'})
+                                {'retry-after': '2'})
         else:
             return super(RetryDriver, self).create_node(name=name,
                     size=size,
@@ -158,7 +164,12 @@ class FakeAwsDriver(FakeDriver):
                                                       auth=auth,
                                                       ex_metadata=ex_metadata,
                                                       ex_userdata=ex_userdata)
-        n.extra = {"launch_time": time.strftime(ARVADOS_TIMEFMT, time.gmtime())[:-1]}
+        n.extra = {
+            "launch_time": time.strftime(ARVADOS_TIMEFMT, time.gmtime())[:-1],
+            "metadata" : {
+                "arvados_node_size": size.id
+            }
+        }
         return n
 
     def list_sizes(self, **kwargs):
@@ -184,7 +195,8 @@ class FakeGceDriver(FakeDriver):
                                                    ex_metadata=ex_metadata)
         n.extra = {
             "metadata": {
-                "items": [{"key": k, "value": v} for k,v in ex_metadata.iteritems()]
+                "items": [{"key": k, "value": v} for k,v in ex_metadata.iteritems()],
+                "arvados_node_size": size.id
             },
             "zone": "fake"
         }