/spark-mongodb/src/test/scala/com/stratio/datasource/mongodb/writer/MongodbWriterIT.scala

https://github.com/Stratio/Spark-MongoDB · Scala · 256 lines · 171 code · 70 blank · 15 comment · 2 complexity · 3334e02c4fcf820209634e8032a580f7 MD5 · raw file

  1. /*
  2. * Copyright (C) 2015 Stratio (http://stratio.com)
  3. *
  4. * Licensed under the Apache License, Version 2.0 (the "License");
  5. * you may not use this file except in compliance with the License.
  6. * You may obtain a copy of the License at
  7. *
  8. * http://www.apache.org/licenses/LICENSE-2.0
  9. *
  10. * Unless required by applicable law or agreed to in writing, software
  11. * distributed under the License is distributed on an "AS IS" BASIS,
  12. * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. * See the License for the specific language governing permissions and
  14. * limitations under the License.
  15. */
  16. package com.stratio.datasource.mongodb.writer
  17. import com.mongodb._
  18. import com.mongodb.casbah.commons.MongoDBObject
  19. import com.mongodb.util.JSON
  20. import com.stratio.datasource.MongodbTestConstants
  21. import com.stratio.datasource.mongodb.client.MongodbClientFactory
  22. import com.stratio.datasource.mongodb.config.{MongodbConfig, MongodbConfigBuilder}
  23. import com.stratio.datasource.mongodb.{MongoEmbedDatabase, TestBsonData}
  24. import org.junit.runner.RunWith
  25. import org.scalatest.junit.JUnitRunner
  26. import org.scalatest.{BeforeAndAfterAll, BeforeAndAfter, FlatSpec, Matchers}
  27. @RunWith(classOf[JUnitRunner])
  28. class MongodbWriterIT extends FlatSpec
  29. with Matchers
  30. with MongoEmbedDatabase
  31. with TestBsonData
  32. with MongodbTestConstants
  33. with BeforeAndAfterAll {
  34. private val host: String = "localhost"
  35. private val collection: String = "testCol"
  36. private val writeConcern = "NORMAL"
  37. private val idField: String = "att2"
  38. private val updateField: Array[String] = Array("att3")
  39. private val wrongIdField: String = "non-existentColumn"
  40. private val language: String = "english"
  41. val testConfig = MongodbConfigBuilder()
  42. .set(MongodbConfig.Host, List(host + ":" + mongoPort))
  43. .set(MongodbConfig.Database, db)
  44. .set(MongodbConfig.Collection, collection)
  45. .set(MongodbConfig.SamplingRatio, 1.0)
  46. .set(MongodbConfig.WriteConcern, writeConcern)
  47. .build()
  48. val testConfigWithPk = MongodbConfigBuilder()
  49. .set(MongodbConfig.Host, List(host + ":" + mongoPort))
  50. .set(MongodbConfig.Database, db)
  51. .set(MongodbConfig.Collection, collection)
  52. .set(MongodbConfig.SamplingRatio, 1.0)
  53. .set(MongodbConfig.WriteConcern, writeConcern)
  54. .build()
  55. val testConfigWithLanguage = MongodbConfigBuilder()
  56. .set(MongodbConfig.Host, List(host + ":" + mongoPort))
  57. .set(MongodbConfig.Database, db)
  58. .set(MongodbConfig.Collection, collection)
  59. .set(MongodbConfig.SamplingRatio, 1.0)
  60. .set(MongodbConfig.WriteConcern, writeConcern)
  61. .set(MongodbConfig.Language, language)
  62. .build()
  63. val testConfigWithWrongPk = MongodbConfigBuilder()
  64. .set(MongodbConfig.Host, List(host + ":" + mongoPort))
  65. .set(MongodbConfig.Database, db)
  66. .set(MongodbConfig.Collection, collection)
  67. .set(MongodbConfig.SamplingRatio, 1.0)
  68. .set(MongodbConfig.WriteConcern, writeConcern)
  69. .build()
  70. val testConfigWithUpdateFields = MongodbConfigBuilder()
  71. .set(MongodbConfig.Host, List(host + ":" + mongoPort))
  72. .set(MongodbConfig.Database, db)
  73. .set(MongodbConfig.Collection, collection)
  74. .set(MongodbConfig.SamplingRatio, 1.0)
  75. .set(MongodbConfig.WriteConcern, writeConcern)
  76. .set(MongodbConfig.UpdateFields, updateField)
  77. .build()
  78. val dbObject = JSON.parse(
  79. """{ "att5" : [ 1 , 2 , 3] ,
  80. "att4" : null ,
  81. "att3" : "hi" ,
  82. "att6" : { "att61" : 1 , "att62" : null } ,
  83. "att2" : 2.0 ,
  84. "att1" : 1}""").asInstanceOf[DBObject]
  85. val listDbObject = List(
  86. JSON.parse(
  87. """{ "att5" : [ 1 , 2 , 3] ,
  88. "att4" : null ,
  89. "att3" : "hi" ,
  90. "att6" : { "att61" : 1 , "att62" : null } ,
  91. "att2" : 2.0 ,
  92. "att1" : 1}""").asInstanceOf[DBObject],
  93. JSON.parse(
  94. """{ "att5" : [ 1 , 2 , 3] ,
  95. "att4" : null ,
  96. "att3" : "holo" ,
  97. "att6" : { "att61" : 1 , "att62" : null } ,
  98. "att2" : 2.0 ,
  99. "att1" : 1}""").asInstanceOf[DBObject])
  100. val updateDbObject = List(
  101. JSON.parse(
  102. """{ "att5" : [ 1 , 2 , 3] ,
  103. "att4" : null ,
  104. "att3" : "holo" ,
  105. "att6" : { "att61" : 1 , "att62" : null } ,
  106. "att2" : 2.0 ,
  107. "att1" : 2}""").asInstanceOf[DBObject])
  108. behavior of "A writer"
  109. it should "properly write in a Mongo collection using the Simple Writer" + scalaBinaryVersion in {
  110. withEmbedMongoFixture(List()) { mongodbProc =>
  111. val mongodbSimpleWriter = new MongodbSimpleWriter(testConfig)
  112. val dbOIterator = List(dbObject).iterator
  113. mongodbSimpleWriter.saveWithPk(dbOIterator)
  114. val mongodbClient = new MongoClient(host, mongoPort)
  115. val dbCollection = mongodbClient.getDB(db).getCollection(collection)
  116. val dbCursor = dbCollection.find()
  117. import scala.collection.JavaConversions._
  118. dbCursor.iterator().toList should equal(List(dbObject))
  119. }
  120. }
  121. it should "properly write in a Mongo collection using the Batch Writer" + scalaBinaryVersion in {
  122. withEmbedMongoFixture(List()) { mongodbProc =>
  123. val mongodbBatchWriter = new MongodbBatchWriter(testConfig)
  124. val dbOIterator = List(dbObject).iterator
  125. mongodbBatchWriter.saveWithPk(dbOIterator)
  126. val mongodbClient = new MongoClient(host, mongoPort)
  127. val dbCollection = mongodbClient.getDB(db).getCollection(collection)
  128. val dbCursor = dbCollection.find()
  129. import scala.collection.JavaConversions._
  130. dbCursor.iterator().toList should equal(List(dbObject))
  131. }
  132. }
  133. it should "manage the incorrect primary key, created in a column that" +
  134. " doesn't exist, rightly" + scalaBinaryVersion in {
  135. withEmbedMongoFixture(List()) { mongodbProc =>
  136. val mongodbBatchWriter = new MongodbBatchWriter(testConfigWithWrongPk)
  137. val dbOIterator = List(dbObject).iterator
  138. mongodbBatchWriter.saveWithPk(dbOIterator)
  139. val mongodbClient = new MongoClient(host, mongoPort)
  140. val dbCollection = mongodbClient.getDB(db).getCollection(collection)
  141. val dbCursor = dbCollection.find()
  142. import scala.collection.JavaConversions._
  143. dbCursor.iterator().toList.forall { case obj: BasicDBObject =>
  144. obj.get("_id") != obj.get("non-existentColumn")
  145. } should be (true)
  146. }
  147. }
  148. it should "manage the language field for text index" + scalaBinaryVersion in {
  149. withEmbedMongoFixture(List()) { mongodbProc =>
  150. val mongodbBatchWriter = new MongodbBatchWriter(testConfigWithLanguage)
  151. val dbOIterator = List(dbObject).iterator
  152. mongodbBatchWriter.saveWithPk(dbOIterator)
  153. val mongodbClient = new MongoClient(host, mongoPort)
  154. val dbCollection = mongodbClient.getDB(db).getCollection(collection)
  155. val dbCursor = dbCollection.find()
  156. import scala.collection.JavaConversions._
  157. dbCursor.iterator().toList.forall { case obj: BasicDBObject =>
  158. obj.get("language") == language
  159. } should be (true)
  160. }
  161. }
  162. it should "manage the update fields and the update query, it has to read the same value from the fields in " +
  163. "configuration" + scalaBinaryVersion in {
  164. withEmbedMongoFixture(List()) { mongodbProc =>
  165. val mongodbBatchWriter = new MongodbBatchWriter(testConfigWithUpdateFields)
  166. val dbOIterator = listDbObject.iterator
  167. val dbUpdateIterator = updateDbObject.iterator
  168. mongodbBatchWriter.saveWithPk(dbOIterator)
  169. val mongodbClient = new MongoClient(host, mongoPort)
  170. val dbCollection = mongodbClient.getDB(db).getCollection(collection)
  171. val dbCursor = dbCollection.find(MongoDBObject("att3" -> "holo"))
  172. import scala.collection.JavaConversions._
  173. dbCursor.iterator().toList.foreach{ case obj: BasicDBObject =>
  174. obj.getInt("att1") should be (1)
  175. }
  176. mongodbBatchWriter.saveWithPk(dbUpdateIterator)
  177. val dbCursor2 = dbCollection.find(MongoDBObject("att3" -> "holo"))
  178. dbCursor2.iterator().toList.foreach { case obj: BasicDBObject =>
  179. obj.getInt("att1") should be (2)
  180. }
  181. }
  182. }
  183. override def afterAll {
  184. MongodbClientFactory.closeAll(false)
  185. }
  186. }