From 203cadf1edb0baa31c65b6111f9bd4f6f3e2203d Mon Sep 17 00:00:00 2001
From: "kirby@puppetlabs.com" <mk.kirby@gmail.com>
Date: Fri, 8 May 2020 15:43:35 -0700
Subject: [PATCH] Add auth from ondemand request to provisioned vms

---
 lib/vmpooler/pool_manager.rb   |  6 ++++--
 spec/helpers.rb                |  4 +++-
 spec/unit/pool_manager_spec.rb | 22 ++++++++++++++++++++++
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/lib/vmpooler/pool_manager.rb b/lib/vmpooler/pool_manager.rb
index e1449f8..0fc5cc1 100644
--- a/lib/vmpooler/pool_manager.rb
+++ b/lib/vmpooler/pool_manager.rb
@@ -138,7 +138,8 @@ module Vmpooler
       finish = format('%<time>.2f', time: Time.now - Time.parse(clone_time))
 
       if request_id
-        if redis.hget("vmpooler__odrequest__#{request_id}", 'status') == 'failed'
+        ondemandrequest_hash = redis.hgetall("vmpooler__odrequest__#{request_id}")
+        if ondemandrequest_hash['status'] == 'failed'
           move_vm_queue(pool, vm, 'pending', 'completed', redis, "moved to completed queue. '#{request_id}' could not be filled in time")
           return nil
         end
@@ -147,9 +148,10 @@ module Vmpooler
         redis.pipelined do
           redis.hset("vmpooler__active__#{pool}", vm, Time.now)
           redis.hset("vmpooler__vm__#{vm}", 'checkout', Time.now)
+          redis.hset("vmpooler__vm__#{vm}", 'token:token', ondemandrequest_hash['token:token']) if ondemandrequest_hash['token:token']
+          redis.hset("vmpooler__vm__#{vm}", 'token:user', ondemandrequest_hash['token:user']) if ondemandrequest_hash['token:user']
           redis.sadd("vmpooler__#{request_id}__#{pool_alias}__#{pool}", vm)
         end
-        puts redis.smembers("vmpooler__#{request_id}__#{pool['alias']}__#{pool}")
         move_vm_queue(pool, vm, 'pending', 'running', redis)
       else
         redis.smove('vmpooler__pending__' + pool, 'vmpooler__ready__' + pool, vm)
diff --git a/spec/helpers.rb b/spec/helpers.rb
index 6547b12..d1019e7 100644
--- a/spec/helpers.rb
+++ b/spec/helpers.rb
@@ -130,9 +130,11 @@ def pool_has_ready_vm?(pool, vm, redis)
   !!redis.sismember('vmpooler__ready__' + pool, vm)
 end
 
-def create_ondemand_request_for_test(request_id, score, platforms_string, redis)
+def create_ondemand_request_for_test(request_id, score, platforms_string, redis, user = nil, token = nil)
   redis.zadd('vmpooler__provisioning__request', score, request_id)
   redis.hset("vmpooler__odrequest__#{request_id}", 'requested', platforms_string)
+  redis.hset("vmpooler__odrequest__#{request_id}", 'token:token', token) if token
+  redis.hset("vmpooler__odrequest__#{request_id}", 'token:user', user) if user
 end
 
 def set_ondemand_request_ready(request_id, redis)
diff --git a/spec/unit/pool_manager_spec.rb b/spec/unit/pool_manager_spec.rb
index 30f451d..192ccc2 100644
--- a/spec/unit/pool_manager_spec.rb
+++ b/spec/unit/pool_manager_spec.rb
@@ -357,6 +357,28 @@ EOT
           end
         end
       end
+
+      context 'with auth on the request' do
+        let(:user) { 'vmpuser' }
+        let(:platform_alias) { pool }
+        let(:platforms_string) { "#{platform_alias}:#{pool}:1" }
+        let(:score) { current_time.to_i }
+        before(:each) do
+          redis_connection_pool.with do |redis|
+            create_ondemand_request_for_test(request_id, score, platforms_string, redis, user, token)
+          end
+        end
+
+        it 'should specify auth data on the vm' do
+          redis_connection_pool.with do |redis|
+            allow(redis).to receive(:hset)
+            expect(redis).to receive(:hset).with("vmpooler__vm__#{vm}", 'token:token', token)
+            expect(redis).to receive(:hset).with("vmpooler__vm__#{vm}", 'token:user', user)
+
+            subject.move_pending_vm_to_ready(vm, pool, redis, request_id)
+          end
+        end
+      end
     end
   end