From 42c675db614dd20812f617181b119af4a7d1b000 Mon Sep 17 00:00:00 2001
From: Stephen MacVicar <Jammjammjamm@users.noreply.github.com>
Date: Fri, 13 Aug 2021 08:25:09 -0400
Subject: [PATCH] Fix for DSTU2/STU3 set_client_on_resource (#144)

* fix setting client on dstu2 and stu3 references

* bump version
---
 lib/fhir_client/client.rb                |  4 ++-
 lib/fhir_client/version.rb               |  2 +-
 test/unit/set_client_on_resource_test.rb | 46 ++++++++++++++++++++++++
 3 files changed, 50 insertions(+), 2 deletions(-)
 create mode 100644 test/unit/set_client_on_resource_test.rb

diff --git a/lib/fhir_client/client.rb b/lib/fhir_client/client.rb
index dff47572..cb61786e 100644
--- a/lib/fhir_client/client.rb
+++ b/lib/fhir_client/client.rb
@@ -349,7 +349,9 @@ def set_client_on_resource(resource)
 
       resource.client = self
       resource.each_element do |element, _, _|
-        element.client = self if element.is_a?(Reference) || element.respond_to?(:resourceType)
+        if element.is_a?(Reference) || element.is_a?(STU3::Reference) || element.is_a?(DSTU2::Reference) || element.respond_to?(:resourceType)
+          element.client = self
+        end
       end
     end
 
diff --git a/lib/fhir_client/version.rb b/lib/fhir_client/version.rb
index a51ff965..e5178052 100644
--- a/lib/fhir_client/version.rb
+++ b/lib/fhir_client/version.rb
@@ -1,5 +1,5 @@
 module FHIR
   class Client
-    VERSION = '5.0.0'.freeze
+    VERSION = '5.0.1'.freeze
   end
 end
diff --git a/test/unit/set_client_on_resource_test.rb b/test/unit/set_client_on_resource_test.rb
new file mode 100644
index 00000000..cd55aee5
--- /dev/null
+++ b/test/unit/set_client_on_resource_test.rb
@@ -0,0 +1,46 @@
+require_relative '../test_helper'
+
+class SetClientOnResourceTest < Test::Unit::TestCase
+  def client
+    @client ||= FHIR::Client.new('abc')
+  end
+
+  def test_r4
+    condition = FHIR::Condition.new(
+      resourceType: 'Condition',
+      subject: {
+        reference: 'Patient/123'
+      }
+    )
+    client.set_client_on_resource(condition)
+
+    assert_equal(client, condition.client)
+    assert_equal(client, condition.subject.client)
+  end
+
+  def test_stu3
+    condition = FHIR::STU3::Condition.new(
+      resourceType: 'Condition',
+      subject: {
+        reference: 'Patient/123'
+      }
+    )
+    client.set_client_on_resource(condition)
+
+    assert_equal(client, condition.client)
+    assert_equal(client, condition.subject.client)
+  end
+
+  def test_dstu2
+    condition = FHIR::DSTU2::Condition.new(
+      resourceType: 'Condition',
+      patient: {
+        reference: 'Patient/123'
+      }
+    )
+    client.set_client_on_resource(condition)
+
+    assert_equal(client, condition.client)
+    assert_equal(client, condition.patient.client)
+  end
+end