diff --git a/lib/resource_registry/versions.rb b/lib/resource_registry/versions.rb index 9724df5..41ec885 100644 --- a/lib/resource_registry/versions.rb +++ b/lib/resource_registry/versions.rb @@ -22,9 +22,14 @@ def find!(name) find(name) || raise("Version '#{name}' not found") end - sig { params(name: String).returns(T.nilable(Version)) } - def find_next(name) - version = find!(name) + sig do + params(name_or_version: T.any(String, Version)).returns( + T.nilable(Version) + ) + end + def find_next(name_or_version) + version = + name_or_version.is_a?(String) ? find!(name_or_version) : name_or_version index = T.must(sorted_versions.index(version)) sorted_versions[index + 1] diff --git a/spec/resource_registry/versions_spec.rb b/spec/resource_registry/versions_spec.rb index b3debb4..d1f16dd 100644 --- a/spec/resource_registry/versions_spec.rb +++ b/spec/resource_registry/versions_spec.rb @@ -51,12 +51,28 @@ it "returns the following version" do expect(subject.find_next("2024-01-01").to_s).to eq("2024-04-01") end + + context "when given a version object" do + it "returns the following version" do + version = subject.find!("2024-01-01") + + expect(subject.find_next(version).to_s).to eq("2024-04-01") + end + end end context "when given the last version" do it "does not returns a version" do expect(subject.find_next("2024-04-01")).to be_nil end + + context "when given a version object" do + it "does not returns a version" do + version = subject.find!("2024-04-01") + + expect(subject.find_next(version)).to be_nil + end + end end context "when given a random value" do