diff --git a/src/main/kotlin/com/test/rest/TestController.kt b/src/main/kotlin/com/test/rest/TestController.kt
index 328cb5a..77f91e9 100644
--- a/src/main/kotlin/com/test/rest/TestController.kt
+++ b/src/main/kotlin/com/test/rest/TestController.kt
@@ -3,6 +3,7 @@ package com.test.rest
 import com.test.model.TestModel
 import com.test.service.TestService
 import org.slf4j.LoggerFactory
+import org.springframework.http.ResponseEntity
 import org.springframework.web.bind.annotation.DeleteMapping
 import org.springframework.web.bind.annotation.GetMapping
 import org.springframework.web.bind.annotation.PathVariable
@@ -22,19 +23,19 @@ open class TestController(private val testService: TestService) {
     //@PreAuthorize("hasRole('USER')")
 
     @PostMapping("/add")
-    fun add(@RequestBody testModel: TestModel): TestModel {
-        return testService.add(testModel)
+    fun add(@RequestBody testModel: TestModel): ResponseEntity<TestModel> {
+        return ResponseEntity.ok(testService.add(testModel))
     }
 
 
     @PutMapping("/update")
-    fun users(@RequestBody testModel: TestModel): TestModel {
-        return testService.update(testModel)
+    fun users(@RequestBody testModel: TestModel): ResponseEntity<TestModel> {
+        return ResponseEntity.ok(testService.update(testModel))
     }
 
     @GetMapping("/{id}")
-    fun get(@PathVariable(value = "id") id: Long): TestModel {
-        return testService.get(id)!!
+    fun get(@PathVariable(value = "id") id: Long): ResponseEntity<TestModel> {
+        return testService.get(id)?.let { ResponseEntity.ok(it) } ?: ResponseEntity.notFound().build()
     }
 
     @DeleteMapping("/{id}")