diff --git a/put_item.go b/put_item.go
index e8f70af..3a33ded 100644
--- a/put_item.go
+++ b/put_item.go
@@ -20,7 +20,8 @@ type PutOption func(*dynamodb.PutItemInput) error
 // Each update increments the version number and if the update fails fetch the record again to get latest version number and try again
 func WithOptimisticLock(key string, currentVersion uint) PutOption {
 	return func(input *dynamodb.PutItemInput) error {
-		condition := "#version = :oldVersion"
+		// Ensure the condition expression is set to check if the version attribute does not exist or matches the old version
+		condition := "attribute_not_exists(#version) or #version = :oldVersion"
 		input.ConditionExpression = &condition
 		if input.ExpressionAttributeNames == nil {
 			input.ExpressionAttributeNames = map[string]string{}
diff --git a/tests/putitem_test.go b/tests/putitem_test.go
index 9b3accf..84bc063 100644
--- a/tests/putitem_test.go
+++ b/tests/putitem_test.go
@@ -109,15 +109,8 @@ type LedgerAccount struct {
 
 func TestPutItemWithOptimisticLock(t *testing.T) {
 	table := prepareTable(t, dynamoEndpoint, "put_optimistic_test")
-	// Create new account item in DynamoDB with default values
-	account := LedgerAccount{ID: "123"}
 	ctx := context.Background()
 	pk := dynago.StringValue("123")
-	err := table.PutItem(ctx, pk, pk, account)
-	if err != nil {
-		t.Fatalf("unexpected error %s", err)
-		return
-	}
 
 	// Update method will add 100 to the current account balance
 	update := func() error {
@@ -151,7 +144,7 @@ func TestPutItemWithOptimisticLock(t *testing.T) {
 	// We expect account balance to be 1000 after 10 update
 	// If any update method overwrote with an outdated value then total balance will be less than 1000
 	var acc LedgerAccount
-	err, _ = table.GetItem(ctx, pk, pk, &acc)
+	err, _ := table.GetItem(ctx, pk, pk, &acc)
 	if err != nil {
 		t.Fatalf("unexpected error %s", err)
 		return